process_ace05ep.py (60914B)
1 """ 2 This script extracts IE annotations from ACE2005 (LDC2006T06). 3 4 Usage: 5 python process_ace.py \ 6 7 """ 8 import os 9 import re 10 import json 11 import glob 12 import tqdm 13 import random 14 import torch 15 from lxml import etree 16 from typing import List, Dict, Any, Tuple 17 from bs4 import BeautifulSoup 18 from dataclasses import dataclass 19 from argparse import ArgumentParser 20 from transformers import (BertTokenizer, 21 RobertaTokenizer, 22 XLMRobertaTokenizer, 23 PreTrainedTokenizer, 24 AutoTokenizer) 25 from nltk import (sent_tokenize as sent_tokenize_, 26 wordpunct_tokenize as wordpunct_tokenize_) 27 #import jieba 28 import stanza 29 30 #nlp_ar = stanza.Pipeline(lang='ar', processors='tokenize') 31 nlp_en = stanza.Pipeline(lang='en', processors='tokenize') 32 #nlp_zh = stanza.Pipeline(lang='zh', processors='tokenize') 33 34 TAG_PATTERN = re.compile('<[^<>]+>', re.MULTILINE) 35 36 DOCS_TO_REVISE_SENT = { 37 'CNN_ENG_20030529_130011.6': [(461, 504), 38 (668, 859), 39 (984, 1074), 40 (1577, 1632)], 41 'CNN_ENG_20030626_203133.11': [(1497, 1527)], 42 'CNN_ENG_20030526_180540.6': [(67, 99)], 43 'CNNHL_ENG_20030523_221118.14': [(136, 174)], 44 'BACONSREBELLION_20050127.1017': [(2659, 2663), 45 (4381, 4405), 46 (410, 458)], 47 'misc.legal.moderated_20050129.2225': [(4118, 4127), 48 (4710, 4794)], 49 'alt.vacation.las-vegas_20050109.0133': [(1201, 1248)], 50 'alt.obituaries_20041121.1339': [(1947, 2044), (1731, 1737)], 51 'APW_ENG_20030326.0190': [(638, 739)], 52 'APW_ENG_20030403.0862': [(729, 781)], 53 'CNN_IP_20030405.1600.02': [(699, 705)], 54 'CNN_IP_20030403.1600.00-1': [(2392, 2399)], 55 'CNN_IP_20030409.1600.04': [(1039, 1050)], 56 'CNN_IP_20030412.1600.03': [(741, 772)], 57 'CNN_IP_20030402.1600.02-1': [(885, 892)], 58 'CNN_IP_20030329.1600.02': [(3229, 3235)], 59 'CNN_IP_20030409.1600.02': [(477, 498)], 60 'CNN_CF_20030304.1900.04': [(522, 575), 61 (5193, 5210), 62 (5461, 5542)], 63 'CNN_IP_20030403.1600.00-3': [(1487, 1493)], 64 'soc.history.war.world-war-ii_20050127.2403': [(414, 441)], 65 'CNN_ENG_20030529_130011.6': [(209, 254), 66 (461, 504), 67 (668, 859), 68 (984, 1074), 69 (1577, 1632)], 70 } 71 72 SKIPPED_DOCS = [ 73 'ALFILFILM_20050202.0740' 74 ] 75 76 def mask_escape(text: str) -> str: 77 """Replaces escaped characters with rare sequences. 78 79 Args: 80 text (str): text to mask. 81 82 Returns: 83 str: masked string. 84 """ 85 return text.replace('&', 'ҪҪҪҪҪ').replace('<', 'ҚҚҚҚ').replace('>', 'ҺҺҺҺ') 86 87 88 def unmask_escape(text: str) -> str: 89 """Replaces masking sequences with the original escaped characters. 90 91 Args: 92 text (str): masked string. 93 94 Returns: 95 str: unmasked string. 96 """ 97 return text.replace('ҪҪҪҪҪ', '&').replace('ҚҚҚҚ', '<').replace('ҺҺҺҺ', '>') 98 99 100 def recover_escape(text: str) -> str: 101 """Converts named character references in the given string to the corresponding 102 Unicode characters. I didn't notice any numeric character references in this 103 dataset. 104 105 Args: 106 text (str): text to unescape. 107 108 Returns: 109 str: unescaped string. 110 """ 111 return text.replace('&', '&').replace('<', '<').replace('>', '>') 112 113 114 def sent_tokenize(text: Tuple[str, int, int], 115 language: str = 'english') -> List[Tuple[str, int, int]]: 116 """ 117 Performs sentence tokenization using stanza 118 119 Args: 120 text (Tuple[str, int, int]): a tuple of three elements, text to split 121 into sentences, start offset, and end offset. 122 language (str): available options: english, chinese, arabic. 123 124 Returns: 125 List[Tuple[str, int, int]]: a list of sentences. 126 """ 127 text, start, end = text 128 if language == 'english': 129 sentences = sent_tokenize_(text, language=language) 130 else: 131 if language == 'chinese': 132 doc = nlp_zh(text) 133 elif language == 'arabic': 134 doc = nlp_ar(text) 135 136 ending_char_idx = [0] 137 for sent in doc.sentences: 138 ending_char_idx.append(sent.tokens[-1].end_char) 139 ending_char_idx = ending_char_idx[:-1] 140 ending_char_idx.append(len(text)) 141 sentences = [] 142 for idx in range(1, len(ending_char_idx)): 143 sentences.append(text[ending_char_idx[idx-1]: ending_char_idx[idx]]) 144 145 last = 0 146 sentences_ = [] 147 for sent in sentences: 148 index = text[last:].find(sent) 149 if index == -1: 150 print(text, sent) 151 else: 152 sentences_.append((sent, last + index + start, 153 last + index + len(sent) + start)) 154 last += index + len(sent) 155 return sentences_ 156 157 158 @dataclass 159 class Span: 160 start: int 161 end: int 162 text: str 163 164 def __post_init__(self): 165 self.start = int(self.start) 166 self.end = int(self.end) 167 self.text = self.text.replace('\n', ' ') 168 169 def char_offsets_to_token_offsets(self, tokens: List[Tuple[int, int, str]]): 170 """Converts self.start and self.end from character offsets to token 171 offsets. 172 173 Args: 174 tokens (List[int, int, str]): a list of token tuples. Each item in 175 the list is a triple (start_offset, end_offset, text). 176 """ 177 start_ = end_ = -1 178 for i, (s, e, _) in enumerate(tokens): 179 if s == self.start: 180 start_ = i 181 if e == self.end: 182 end_ = i + 1 183 if start_ == -1 or end_ == -1 or start_ > end_: 184 raise ValueError('Failed to update offsets for {}-{}:{} in {}'.format( 185 self.start, self.end, self.text, tokens)) 186 self.start, self.end = start_, end_ 187 188 def to_dict(self) -> Dict[str, Any]: 189 """Converts instance variables to a dict. 190 191 Returns: 192 dict: a dict of instance variables. 193 """ 194 return { 195 'text': recover_escape(self.text), 196 'start': self.start, 197 'end': self.end, 198 } 199 200 def remove_space(self): 201 """Removes heading and trailing spaces in the span text.""" 202 # heading spaces 203 text = self.text.lstrip(' ') 204 self.start += len(self.text) - len(text) 205 # trailing spaces 206 after_text = text.rstrip(' ') 207 #self.end = self.start + len(text) 208 #self.text = text 209 self.end += len(after_text) - len(text) 210 self.text = after_text 211 212 213 def copy(self): 214 """Makes a copy of itself. 215 216 Returns: 217 Span: a copy of itself.""" 218 return Span(self.start, self.end, self.text) 219 220 221 @dataclass 222 class Entity(Span): 223 entity_id: str 224 mention_id: str 225 entity_type: str 226 entity_subtype: str 227 mention_type: str 228 value: str = None 229 230 def copy(self): 231 return Entity(self.start, self.end, self.text, 232 self.entity_id, self.mention_id, self.entity_type, 233 self.entity_subtype, self.mention_type, self.value) 234 235 def to_dict(self) -> Dict[str, Any]: 236 """Converts instance variables to a dict. 237 238 Returns: 239 Dict: a dict of instance variables. 240 """ 241 entity_dict = { 242 'text': recover_escape(self.text), 243 'entity_id': self.entity_id, 244 'mention_id': self.mention_id, 245 'start': self.start, 246 'end': self.end, 247 'entity_type': self.entity_type, 248 'entity_subtype': self.entity_subtype, 249 'mention_type': self.mention_type, 250 } 251 if self.value: 252 entity_dict['value'] = self.value 253 return entity_dict 254 255 @dataclass 256 class Entity_whole: 257 entity_id: str 258 entity_type: str 259 entity_subtype: str 260 entity_mentions: List[Entity] 261 262 def to_dict(self) -> Dict[str, Any]: 263 """Converts instance variables to a dict. 264 265 Returns: 266 Dict: a dict of instance variables. 267 """ 268 entity_dict = { 269 'entity_id': self.entity_id, 270 'entity_type': self.entity_type, 271 'entity_subtype': self.entity_subtype, 272 'entity_mentions': [e.to_dict() for e in self.entity_mentions] 273 } 274 return entity_dict 275 276 def copy(self): 277 """Makes a copy of itself. 278 """ 279 return Entity_whole(self.entity_id, self.entity_type, 280 self.entity_subtype, self.entity_mentions) 281 282 283 @dataclass 284 class RelationArgument: 285 mention_id: str 286 role: str 287 text: str 288 289 def to_dict(self) -> Dict[str, Any]: 290 """Converts instance variables to a dict. 291 292 Returns: 293 Dict[str, Any]: a dict of instance variables. 294 """ 295 return { 296 'mention_id': self.mention_id, 297 'role': self.role, 298 'text': recover_escape(self.text) 299 } 300 301 @dataclass 302 class RelEntBaseArgument: 303 entity_id: str 304 role: str 305 306 def to_dict(self) -> Dict[str, Any]: 307 """Converts instance variables to a dict. 308 309 Returns: 310 Dict[str, Any]: a dict of instance variables. 311 """ 312 return { 313 'entity_id': self.entity_id, 314 'role': self.role, 315 } 316 317 @dataclass 318 class Relation: 319 relation_id: str 320 relation_type: str 321 relation_subtype: str 322 arg1: RelationArgument 323 arg2: RelationArgument 324 325 def to_dict(self) -> Dict[str, Any]: 326 """Converts instance variables to a dict. 327 328 Returns: 329 Dict[str, Any]: a dict of instance variables. 330 """ 331 return { 332 'relation_id': self.relation_id, 333 'relation_type': self.relation_type, 334 'relation_subtype': self.relation_subtype, 335 'arg1': self.arg1.to_dict(), 336 'arg2': self.arg2.to_dict(), 337 } 338 339 @dataclass 340 class RelEntBase: 341 relation_id: str 342 relation_type: str 343 relation_subtype: str 344 arg1: RelEntBaseArgument 345 arg2: RelEntBaseArgument 346 347 def to_dict(self) -> Dict[str, Any]: 348 """Converts instance variables to a dict. 349 350 Returns: 351 Dict[str, Any]: a dict of instance variables. 352 """ 353 return { 354 'relation_id': self.relation_id, 355 'relation_type': self.relation_type, 356 'relation_subtype': self.relation_subtype, 357 'arg1': self.arg1.to_dict(), 358 'arg2': self.arg2.to_dict(), 359 } 360 361 @dataclass 362 class EventArgument: 363 mention_id: str 364 role: str 365 text: str 366 367 def to_dict(self) -> Dict[str, Any]: 368 """Converts instance variables to a dict. 369 370 Returns: 371 Dict[str, Any]: a dict of instance variables. 372 """ 373 return { 374 'mention_id': self.mention_id, 375 'role': self.role, 376 'text': recover_escape(self.text), 377 } 378 379 @dataclass 380 class Event: 381 event_id: str 382 mention_id: str 383 event_type: str 384 event_subtype: str 385 trigger: Span 386 arguments: List[EventArgument] 387 388 def copy(self): 389 return Event(self.event_id, self.mention_id, self.event_type, 390 self.event_subtype, self.trigger.copy(), self.arguments) 391 392 def to_dict(self) -> Dict[str, Any]: 393 """Converts instance variables to a dict. 394 395 Returns: 396 Dict[str, Any]: a dict of instance variables. 397 """ 398 return { 399 'event_id': self.event_id, 400 'mention_id': self.mention_id, 401 'event_type': self.event_type, 402 'event_subtype': self.event_subtype, 403 'trigger': self.trigger.to_dict(), 404 'arguments': [arg.to_dict() for arg in self.arguments], 405 } 406 407 @dataclass 408 class EventArgument_whole: 409 entity_id: str 410 role: str 411 412 def to_dict(self) -> Dict[str, Any]: 413 """Converts instance variables to a dict. 414 415 Returns: 416 Dict[str, Any]: a dict of instance variables. 417 """ 418 return { 419 'entity_id': self.entity_id, 420 'role': self.role, 421 } 422 423 @dataclass 424 class Event_whole: 425 event_id: str 426 event_type: str 427 event_subtype: str 428 arguments: List[EventArgument_whole] 429 event_mentions: List[Event] 430 431 def to_dict(self) -> Dict[str, Any]: 432 """Converts instance variables to a dict. 433 434 Returns: 435 Dict[str, Any]: a dict of instance variables. 436 """ 437 return { 438 'event_id': self.event_id, 439 'event_type': self.event_type, 440 'event_subtype': self.event_subtype, 441 'arguments': [arg.to_dict() for arg in self.arguments], 442 'event_mentions': [e.to_dict() for e in self.event_mentions] 443 } 444 445 446 @dataclass 447 class Sentence(Span): 448 sent_id: str 449 tokens: List[str] 450 sent_starts: List[int] 451 entities: List[Entity] 452 relations: List[Relation] 453 events: List[Event] 454 entity_cluster: List[Entity_whole] 455 relations_entbase: List[RelEntBase] 456 event_cluster: List[Event_whole] 457 458 def to_dict(self) -> Dict[str, Any]: 459 """Converts instance variables to a dict. 460 461 Returns: 462 Dict[str, Any]: a dict of instance variables. 463 """ 464 return { 465 'sent_id': self.sent_id, 466 'tokens': [recover_escape(t) for t in self.tokens], 467 'sentence_starts': self.sent_starts, 468 'entities': [entity.to_dict() for entity in self.entities], 469 'relations': [relation.to_dict() for relation in self.relations], 470 'events': [event.to_dict() for event in self.events], 471 'entity_cluster': [entity.to_dict() for entity in self.entity_cluster], 472 'event_cluster': [event.to_dict() for event in self.event_cluster], 473 'relations_entbase': [relation.to_dict() for relation in self.relations_entbase], 474 'start': self.start, 475 'end': self.end, 476 'text': recover_escape(self.text).replace('\t', ' '), 477 } 478 479 480 @dataclass 481 class Document: 482 doc_id: str 483 sentences: List[Sentence] 484 485 def to_dict(self) -> Dict[str, Any]: 486 """Converts instance variables to a dict. 487 488 Returns: 489 Dict[str, Any]: a dict of instance variables. 490 """ 491 return { 492 'doc_id': self.doc_id, 493 'sentences': [sent.to_dict() for sent in self.sentences] 494 } 495 496 497 def revise_sentences(sentences: List[Tuple[str, int, int]], 498 doc_id: str) -> List[Tuple[int, int, str]]: 499 """Automatic sentence tokenization may have errors for a few documents. 500 501 Args: 502 sentences (List[Tuple[str, int, int]]): a list of sentence tuples. 503 doc_id (str): document ID. 504 505 Returns: 506 List[Tuple[str, int, int]]: a list of revised sentence tuples. 507 """ 508 sentences_ = [] 509 510 offset_list = DOCS_TO_REVISE_SENT[doc_id] 511 first_part_offsets = {offset for offset, _ in offset_list} 512 second_part_offsets = {offset for _, offset in offset_list} 513 514 515 for sentence_idx, (text, start, end) in enumerate(sentences): 516 if start in first_part_offsets: 517 next_text, next_start, next_end = sentences[sentence_idx + 1] 518 space = ' ' * (next_start - end) 519 sentences_.append((text + space + next_text, start, next_end)) 520 elif start in second_part_offsets: 521 continue 522 else: 523 sentences_.append((text, start, end)) 524 525 return sentences_ 526 527 528 def read_sgm_file(path: str, 529 language: str = 'english') -> List[Tuple[str, int, int]]: 530 """Reads a SGM text file. 531 532 Args: 533 path (str): path to the input file. 534 language (str): document language. Valid values: "english", "arabic" or "chinese". 535 536 Returns: 537 List[Tuple[str, int, int]]: a list of sentences. Each item in the list 538 is a tuple of three elements, sentence text, start offset, and end 539 offset. 540 """ 541 data = open(path, 'r', encoding='utf-8').read() 542 543 # Chunk the document 544 chunks = TAG_PATTERN.sub('⁑', data).split('⁑') 545 546 # Get the offset of <TEXT> 547 data = data.replace('<TEXT>', '⁂') 548 data = TAG_PATTERN.sub('', data) 549 min_offset = max(0, data.find('⁂')) 550 data = data.replace('⁂', '') 551 552 # Extract sentences from chunks 553 chunk_offset = 0 554 sentences = [] 555 for chunk in chunks: 556 lines = chunk.split('\n') 557 current_sentence = [] 558 start = offset = 0 559 for line in lines: 560 offset += len(line) + 1 561 if line.strip(): 562 current_sentence.append(line) 563 else: 564 # empty line 565 if current_sentence: 566 sentence = ' '.join(current_sentence) 567 if start + chunk_offset >= min_offset: 568 sentences.append((sentence, 569 start + chunk_offset, 570 start + chunk_offset + len(sentence))) 571 current_sentence = [] 572 start = offset 573 if current_sentence: 574 sentence = ' '.join(current_sentence) 575 if start + chunk_offset >= min_offset: 576 sentences.append((sentence, 577 start + chunk_offset, 578 start + chunk_offset + len(sentence))) 579 chunk_offset += len(chunk) 580 581 # Re-tokenize sentences 582 sentences = [s for sent in sentences 583 for s in sent_tokenize(sent, language=language)] 584 585 return sentences 586 587 def read_apf_file(path: str, 588 time_and_val: bool = False 589 ) -> Tuple[str, str, List[Entity], List[Relation], List[Event]]: 590 """Reads an APF file. 591 592 Args: 593 path (str): path to the input file. 594 time_and_val (bool): extract times and values or not. 595 596 Returns: 597 doc_id (str): document ID. 598 source (str): document source. 599 entity_list (List[Entity]): a list of Entity instances. 600 relation_list (List[Relation]): a list of Relation instances. 601 event_list (List[Event]): a list of Events instances. 602 """ 603 data = open(path, 'r', encoding='utf-8').read() 604 soup = BeautifulSoup(data, 'lxml-xml') 605 606 # metadata 607 root = soup.find('source_file') 608 source = root['SOURCE'] 609 doc = root.find('document') 610 doc_id = doc['DOCID'] 611 612 entity_list, relation_list, event_list = [], [], [] 613 entity_clusters, event_clusters, relentbase_list = [], [], [] 614 615 616 # entities: nam, nom, pro 617 for entity in doc.find_all('entity'): 618 entity_id = entity['ID'] 619 entity_type = entity['TYPE'] 620 entity_subtype = entity['SUBTYPE'] 621 entity_cluster = [] 622 for entity_mention in entity.find_all('entity_mention'): 623 mention_id = entity_mention['ID'] 624 mention_type = entity_mention['TYPE'] 625 head = entity_mention.find('head').find('charseq') 626 start, end, text = int(head['START']), int(head['END']), head.text 627 entity_list.append(Entity(start, end+1, text, 628 entity_id, mention_id, entity_type, 629 entity_subtype, mention_type)) 630 entity_cluster.append(Entity(start, end+1, text, 631 entity_id, mention_id, entity_type, 632 entity_subtype, mention_type)) 633 entity_clusters.append(Entity_whole(entity_id, entity_type, entity_subtype, entity_cluster)) 634 if time_and_val: 635 # entities: value 636 for entity in doc.find_all('value'): 637 enitty_id = entity['ID'] 638 entity_type = entity['TYPE'] 639 entity_subtype = entity.get('SUBTYPE', None) 640 entity_cluster = [] 641 for entity_mention in entity.find_all('value_mention'): 642 mention_id = entity_mention['ID'] 643 mention_type = 'VALUE' 644 extent = entity_mention.find('extent').find('charseq') 645 start, end, text = int(extent['START']), int( 646 extent['END']), extent.text 647 entity_list.append(Entity(start, end+1, text, 648 entity_id, mention_id, entity_type, 649 entity_subtype, mention_type)) 650 entity_cluster.append(Entity(start, end+1, text, 651 entity_id, mention_id, entity_type, 652 entity_subtype, mention_type)) 653 entity_clusters.append(Entity_whole(entity_id, entity_type, entity_subtype, entity_cluster)) 654 655 # entities: timex 656 for entity in doc.find_all('timex2'): 657 entity_id = entity['ID'] 658 enitty_type = entity_subtype = 'TIME' 659 value = entity.get('VAL', None) 660 entity_cluster = [] 661 for entity_mention in entity.find_all('timex2_mention'): 662 mention_id = entity_mention['ID'] 663 mention_type = 'TIME' 664 extent = entity_mention.find('extent').find('charseq') 665 start, end, text = int(extent['START']), int( 666 extent['END']), extent.text 667 entity_list.append(Entity(start, end+1, text, 668 entity_id, mention_id, entity_type, 669 entity_subtype, mention_type, 670 value=value)) 671 entity_cluster.append(Entity(start, end+1, text, 672 entity_id, mention_id, entity_type, 673 entity_subtype, mention_type, 674 value=value)) 675 entity_clusters.append(Entity_whole(entity_id, entity_type, entity_subtype, entity_cluster)) 676 677 # relations 678 for relation in doc.find_all('relation'): 679 relation_id = relation['ID'] 680 relation_type = relation['TYPE'] 681 if relation_type == 'METONYMY': 682 continue 683 relation_subtype = relation['SUBTYPE'] 684 for relation_mention in relation.find_all('relation_mention'): 685 mention_id = relation_mention['ID'] 686 arg1 = arg2 = None 687 for arg in relation_mention.find_all('relation_mention_argument'): 688 arg_mention_id = arg['REFID'] 689 arg_role = arg['ROLE'] 690 arg_text = arg.find('extent').find('charseq').text 691 if arg_role == 'Arg-1': 692 arg1 = RelationArgument(arg_mention_id, arg_role, arg_text) 693 elif arg_role == 'Arg-2': 694 arg2 = RelationArgument(arg_mention_id, arg_role, arg_text) 695 if arg1 and arg2: 696 relation_list.append(Relation(mention_id, relation_type, 697 relation_subtype, arg1, arg2)) 698 arg1 = arg2 = None 699 for relation_argument in relation.find_all('relation_argument'): 700 arg_id = arg['REFID'] 701 arg_role = arg['ROLE'] 702 if arg_role == 'Arg-1': 703 arg1 = RelEntBaseArgument(arg_id, arg_role) 704 elif arg_role == 'Arg-2': 705 arg2 = RelEntBaseArgument(arg_id, arg_role) 706 if arg1 and arg2: 707 relentbase_list.append(RelEntBase(relation_id, relation_type, relation_subtype, 708 arg1, arg2)) 709 710 # events 711 for event in doc.find_all('event'): 712 event_id = event['ID'] 713 event_type = event['TYPE'] 714 event_subtype = event['SUBTYPE'] 715 event_modality = event['MODALITY'] 716 event_polarity = event['POLARITY'] 717 event_genericity = event['GENERICITY'] 718 event_tense = event['TENSE'] 719 event_cluster = [] 720 for event_mention in event.find_all('event_mention'): 721 mention_id = event_mention['ID'] 722 trigger = event_mention.find('anchor').find('charseq') 723 trigger_start, trigger_end = int( 724 trigger['START']), int(trigger['END']) 725 trigger_text = trigger.text 726 event_args = [] 727 for arg in event_mention.find_all('event_mention_argument'): 728 arg_mention_id = arg['REFID'] 729 arg_role = arg['ROLE'] 730 arg_text = arg.find('extent').find('charseq').text 731 event_args.append(EventArgument( 732 arg_mention_id, arg_role, arg_text)) 733 event_list.append(Event(event_id, mention_id, 734 event_type, event_subtype, 735 Span(trigger_start, 736 trigger_end + 1, trigger_text), 737 event_args)) 738 event_cluster.append(Event(event_id, mention_id, 739 event_type, event_subtype, 740 Span(trigger_start, 741 trigger_end + 1, trigger_text), 742 event_args)) 743 event_arguments = [] 744 for arg in event.find_all('event_argument'): 745 arg_id = arg['REFID'] 746 arg_role = arg['ROLE'] 747 event_arguments.append(EventArgument_whole(arg_id, arg_role)) 748 event_clusters.append(Event_whole(event_id, event_type, event_subtype, 749 event_arguments, event_cluster)) 750 751 # remove heading/tailing spaces 752 for entity in entity_list: 753 entity.remove_space() 754 for event in event_list: 755 event.trigger.remove_space() 756 for entities in entity_clusters: 757 for entity in entities.entity_mentions: 758 entity.remove_space() 759 for events in event_clusters: 760 for event in events.event_mentions: 761 event.trigger.remove_space() 762 763 return (doc_id, source, entity_list, relation_list, event_list, 764 entity_clusters, event_clusters, relentbase_list) 765 766 767 def process_entities(entities: List[Entity], 768 sentences: List[List[Tuple[str, int, int]]] 769 ) -> List[List[Entity]]: 770 """Cleans entities and splits them into lists 771 772 Args: 773 entities (List[Entity]): a list of Entity instances. 774 sentences (List[List[Tuple[str, int, int]]]): a list of window of sentences. 775 776 Returns: 777 List[List[Entity]]: a list of sentence entity lists. 778 """ 779 sentence_entities = [[] for _ in range(len(sentences))] 780 781 # assign each entity to the sentence where it appears 782 for entity in entities: 783 start, end = entity.start, entity.end 784 for i, sentence in enumerate(sentences): # each window 785 flag = False 786 for text, s, e in sentence: 787 if start >= s and end <= e: 788 # put this entity into this window 789 flag = True 790 break 791 if flag: 792 sentence_entities[i].append(entity.copy()) 793 # TODO: I'm not sure whether there will be some entities that 794 # cross single sentence boundary 795 796 """ 797 # remove overlapping entities 798 # TODO: if you want to use it, you need to modify some in order to fit "window" setting 799 sentence_entities_cleaned = [[] for _ in range(len(sentences))] 800 for i, entities in enumerate(sentence_entities): 801 if not entities: 802 continue 803 # prefer longer entities 804 entities.sort(key=lambda x: (x.end - x.start), reverse=True) 805 chars = [0] * max([x.end for x in entities]) 806 for entity in entities: 807 overlap = False 808 for j in range(entity.start, entity.end): 809 if chars[j] == 1: 810 overlap = True 811 break 812 if not overlap: 813 chars[entity.start:entity.end] = [ 814 1] * (entity.end - entity.start) 815 sentence_entities_cleaned[i].append(entity) 816 sentence_entities_cleaned[i].sort(key=lambda x: x.start) 817 """ 818 return sentence_entities 819 820 821 def process_entity_cluster(entity_cluster: List[Entity_whole], 822 sentences: List[List[Tuple[str, int, int]]] 823 ) -> List[List[Entity_whole]]: 824 """Cleans entity_cluster and splits them into lists 825 826 Args: 827 entities (List[Entity_whole]): a list of Entity_whole instances. 828 sentences (List[List[Tuple[str, int, int]]]): a list of window of sentences. 829 830 Returns: 831 List[List[Entity_whole]]: a list of sentence Entity_whole lists. 832 """ 833 sentence_entities = [[] for _ in range(len(sentences))] 834 835 # assign each entity_whole to the sentence where it appears 836 for i, sentence in enumerate(sentences): # each window of sentence 837 for entity in entity_cluster: 838 new_mention_list = [] 839 for mention in entity.entity_mentions: 840 start, end = mention.start, mention.end 841 for _, s, e in sentence: 842 if start >= s and end <= e: 843 # it means we can find this mention in this window 844 new_mention_list.append(mention.copy()) 845 break 846 if any(new_mention_list): 847 new_entity_whole = entity.copy() 848 new_entity_whole.entity_mentions = new_mention_list 849 sentence_entities[i].append(new_entity_whole) 850 return sentence_entities 851 852 853 def process_events(events: List[Event], 854 sentence_entities: List[List[Entity]], 855 sentences: List[List[Tuple[str, int, int]]] 856 ) -> List[List[Event]]: 857 """Cleans and assigns events. 858 859 Args: 860 events (List[Event]): A list of Event objects 861 entence_entities (List[List[Entity]]): A list of sentence entity lists. 862 sentences (List[List[Tuple[str, int, int]]]): a list of window of sentences. 863 864 Returns: 865 List[List[Event]]: a list of sentence event lists. 866 """ 867 sentence_events = [[] for _ in range(len(sentences))] 868 # assign each event mention to the sentence where it appears 869 for event in events: 870 start, end = event.trigger.start, event.trigger.end 871 for i, sentence in enumerate(sentences): # each window 872 # check trigger 873 flag = False 874 for _, s, e in sentence: # every sentence in this window 875 if start >= s and end <= e: 876 flag = True 877 break 878 if flag: 879 sent_entities = sentence_entities[i] 880 # clean the argument list 881 arguments = [] 882 for argument in event.arguments: 883 mention_id = argument.mention_id 884 for entity in sent_entities: 885 if entity.mention_id == mention_id: 886 arguments.append(argument) 887 event_cleaned = Event(event.event_id, event.mention_id, 888 event.event_type, event.event_subtype, 889 trigger=event.trigger.copy(), 890 arguments=arguments) 891 sentence_events[i].append(event_cleaned) 892 """ 893 # remove overlapping events 894 # TODO: if you want to use it, you need to modify some in order to fit "window" setting 895 sentence_events_cleaned = [[] for _ in range(len(sentences))] 896 for i, events in enumerate(sentence_events): 897 if not events: 898 continue 899 events.sort(key=lambda x: (x.trigger.end - x.trigger.start), 900 reverse=True) 901 chars = [0] * max([x.trigger.end for x in events]) 902 for event in events: 903 overlap = False 904 for j in range(event.trigger.start, event.trigger.end): 905 if chars[j] == 1: 906 overlap = True 907 break 908 if not overlap: 909 chars[event.trigger.start:event.trigger.end] = [ 910 1] * (event.trigger.end - event.trigger.start) 911 sentence_events_cleaned[i].append(event) 912 sentence_events_cleaned[i].sort(key=lambda x: x.trigger.start) 913 """ 914 return sentence_events 915 916 917 def process_event_cluster(event_clusters: List[Event_whole], 918 sentence_events: List[List[Event]], 919 sentences: List[List[Tuple[str, int, int]]] 920 ) -> List[List[Event_whole]]: 921 """Cleans and assigns event cluster. 922 923 Args: 924 event_clusters (List[Event_whole]): A list of event cluster objects 925 sentence_events (List[List[Event]]): A list of sentence event lists. 926 sentences (List[List[Tuple[str, int, int]]]): a list of window of sentences. 927 928 Returns: 929 List[List[Event_whole]]: a list of sentence event cluster lists. 930 """ 931 sentence_events_cluster = [[] for _ in range(len(sentences))] 932 # Reconstruct event_whole based on the event mentions that within the sentence(window) span 933 for event_cluster in event_clusters: 934 for i, sentence_event in enumerate(sentence_events): 935 within_events = [eve.copy() for eve in sentence_event if (eve.event_id == event_cluster.event_id)] 936 if any(within_events): 937 arguments = [] 938 for argu in event_cluster.arguments: 939 flag = False 940 for within_event in within_events: 941 for arg in within_event.arguments: 942 if arg.mention_id.rsplit('-', 1)[0] == argu.entity_id: 943 flag = True 944 break 945 if flag: 946 break 947 if flag: 948 arguments.append(argu) 949 sentence_events_cluster[i].append(Event_whole(event_cluster.event_id, event_cluster.event_type, 950 event_cluster.event_subtype, arguments, within_events)) 951 952 return sentence_events_cluster 953 954 955 def process_relation(relations: List[Relation], 956 sentence_entities: List[List[Entity]], 957 sentences: List[List[Tuple[str, int, int]]] 958 ) -> List[List[Relation]]: 959 """Cleans and assigns relations 960 961 Args: 962 relations (List[Relation]): a list of Relation instances. 963 sentence_entities (List[List[Entity]]): a list of sentence entity lists. 964 sentences (List[List[Tuple[str, int, int]]]): a list of window of sentences. 965 966 Returns: 967 List[List[Relation]]: a list of sentence relation lists. 968 """ 969 sentence_relations = [[] for _ in range(len(sentences))] 970 for relation in relations: 971 mention_id1 = relation.arg1.mention_id 972 mention_id2 = relation.arg2.mention_id 973 for i, entities in enumerate(sentence_entities): 974 arg1_in_sent = any([mention_id1 == e.mention_id for e in entities]) 975 arg2_in_sent = any([mention_id2 == e.mention_id for e in entities]) 976 if arg1_in_sent and arg2_in_sent: 977 sentence_relations[i].append(relation) 978 return sentence_relations 979 980 981 def process_relation_entbase(relations: List[RelEntBase], 982 sentence_relations: List[List[Relation]], 983 sentences: List[List[Tuple[str, int, int]]] 984 )-> List[List[RelEntBase]]: 985 sentence_relations_entbase = [[] for _ in range(len(sentences))] 986 for relation in relations: 987 for i, sentence_relation in enumerate(sentence_relations): # relation within the window 988 # if we can find relation mentions within the span, we add this entity-based relation 989 if any([rel_mention for rel_mention in sentence_relation if ((rel_mention.relation_id).rsplit('-',1)[0] == relation.relation_id)]): 990 sentence_relations_entbase[i].append(relation) 991 else: 992 for rel_mention in sentence_relations: 993 print(rel_mention.relation_id.rsplit('-',1)[0]) 994 print(relation.relation_id) 995 return sentence_relations_entbase 996 997 998 def tokenize(sentences: List[Tuple[str, int, int]], 999 entities: List[Entity], 1000 events: List[Event], 1001 language: str = 'english' 1002 ) -> List[Tuple[int, int, str]]: 1003 """Tokenizes a sentence. 1004 Each sentence is first split into chunks that are entity/event spans or words 1005 between two spans. After that, word tokenization is performed on each chunk. 1006 1007 Args: 1008 sentences (List[Tuple[str, int, int]]): a window of sentences (tuple (text, start, end)). 1009 entities (List[Entity]): A list of Entity instances. 1010 events (List[Event]): A list of Event instances. 1011 1012 Returns: 1013 List[Tuple[int, int, str]]: a list of token tuples. Each tuple consists 1014 of three elements, start offset, end offset, and token text. 1015 """ 1016 all_tokens = [] 1017 token_starts = [0] 1018 for sentence in sentences: 1019 text, start, end= sentence 1020 text = mask_escape(text) 1021 1022 # split the sentence into chunks 1023 splits = {0, len(text)} 1024 for entity in entities: 1025 if entity.start >= start and entity.end <= end: 1026 splits.add(entity.start - start) 1027 splits.add(entity.end - start) 1028 for event in events: 1029 if event.trigger.start >= start and event.trigger.end <= end: 1030 splits.add(event.trigger.start - start) 1031 splits.add(event.trigger.end - start) 1032 splits = sorted(list(splits)) 1033 chunks = [(splits[i], splits[i + 1], text[splits[i]:splits[i + 1]]) 1034 for i in range(len(splits) - 1)] 1035 1036 tokens = [] 1037 if language == 'chinese': 1038 def _tokenize_chinese(text): 1039 return [c for c in jieba.cut(text) if c.strip()] 1040 1041 # tokenize each chunk 1042 chunks = [(s, e, t, _tokenize_chinese(t)) 1043 for s, e, t in chunks] 1044 1045 # merge chunks and add word offsets 1046 for chunk_start, chunk_end, chunk_text, chunk_tokens in chunks: 1047 last = 0 1048 chunk_tokens_ = [] 1049 for token in chunk_tokens: 1050 token_start = chunk_text[last:].find(token) 1051 if token_start == -1: 1052 raise ValueError( 1053 'Cannot find token {} in {}'.format(token, text)) 1054 token_end = token_start + len(token) 1055 chunk_tokens_.append((token_start + start + last + chunk_start, 1056 token_end + start + last + chunk_start, 1057 unmask_escape(token))) 1058 last += token_end 1059 tokens.extend(chunk_tokens_) 1060 #else: 1061 elif language == 'arabic': 1062 for chunk_start, _, t in chunks: 1063 if language == 'english': 1064 doc = nlp_en(t) 1065 elif language == 'arabic': 1066 doc = nlp_ar(t) 1067 for sent in doc.sentences: 1068 for tok in sent.tokens: 1069 if tok.text != '': 1070 tokens.append((start + chunk_start + tok.start_char, 1071 start + chunk_start + tok.end_char, 1072 unmask_escape(tok.text))) 1073 elif language == 'english': 1074 # TODO: In order to make a fair comparison with OneIE, we use NLTK's tokenizer rather than Stanza for English 1075 # tokenize each chunk 1076 chunks = [(s, e, t, wordpunct_tokenize_(t)) 1077 for s, e, t in chunks] 1078 1079 # merge chunks and add word offsets 1080 tokens = [] 1081 for chunk_start, chunk_end, chunk_text, chunk_tokens in chunks: 1082 last = 0 1083 chunk_tokens_ = [] 1084 for token in chunk_tokens: 1085 token_start = chunk_text[last:].find(token) 1086 if token_start == -1: 1087 raise ValueError( 1088 'Cannot find token {} in {}'.format(token, text)) 1089 token_end = token_start + len(token) 1090 chunk_tokens_.append((token_start + start + last + chunk_start, 1091 token_end + start + last + chunk_start, 1092 unmask_escape(token))) 1093 last += token_end 1094 tokens.extend(chunk_tokens_) 1095 1096 all_tokens.extend(tokens) 1097 token_starts.append(len(all_tokens)) 1098 1099 return all_tokens, token_starts 1100 1101 def sentence2window(sentences, window_size_): 1102 """Convert sentences list to windows 1103 1104 Args: 1105 sentences (List[Tuple[str, int, int]]) 1106 window_size : int 1107 1108 Returns: 1109 List[List[Tuple[str, int, int]]] 1110 1111 """ 1112 if window_size_ > len(sentences): 1113 window_size = len(sentences) 1114 else: 1115 window_size = window_size_ 1116 1117 output = [] 1118 for i in range(len(sentences)-window_size+1): 1119 output.append(sentences[i:i+window_size]) 1120 return output 1121 1122 def convert(sgm_file: str, 1123 apf_file: str, 1124 time_and_val: bool = False, 1125 language: str = 'english', 1126 window_size: int = 1) -> Document: 1127 """Converts a document. 1128 1129 Args: 1130 sgm_file (str): path to a SGM file. 1131 apf_file (str): path to a APF file. 1132 time_and_val (bool, optional): extracts times and values or not. 1133 Defaults to False. 1134 language (str, optional): document language. Available options: english, 1135 chinese, chinese. Defaults to 'english'. 1136 1137 Returns: 1138 Document: a Document instance. 1139 """ 1140 sentences = read_sgm_file(sgm_file, language=language) 1141 doc_id, source, entities, relations, events, entity_cluster, event_cluster, relentbase =\ 1142 read_apf_file(apf_file, time_and_val=time_and_val) 1143 1144 if doc_id in SKIPPED_DOCS: 1145 return False 1146 1147 # Reivse sentences 1148 if doc_id in DOCS_TO_REVISE_SENT: 1149 sentences = revise_sentences(sentences, doc_id) 1150 1151 # concate windows 1152 sentences = sentence2window(sentences, window_size) 1153 1154 # Process entities, relations, and events 1155 sentence_entities = process_entities(entities, sentences) 1156 sentence_relations = process_relation( 1157 relations, sentence_entities, sentences) 1158 sentence_events = process_events(events, sentence_entities, sentences) 1159 1160 # Process entity_cluster, event_cluster, relation(entity_base) 1161 sentence_entity_cluster = process_entity_cluster(entity_cluster, sentences) 1162 sentence_relentbase = process_relation_entbase( 1163 relentbase, sentence_relations, sentences) 1164 sentence_event_cluster = process_event_cluster(event_cluster, sentence_events, sentences) 1165 1166 # Tokenization 1167 sentence_tokens = [] 1168 sentence_starts = [] 1169 for s, ent, evt in zip(sentences, sentence_entities, sentence_events): 1170 sent_tokens, sent_start = tokenize(s, ent, evt, language=language) 1171 sentence_tokens.append(sent_tokens) 1172 sentence_starts.append(sent_start) 1173 1174 # Convert span character offsets to token indices 1175 sentence_objs = [] 1176 for i, (toks, starts, ents, evts, rels, ent_cls, evt_cls, rel_ents, sent) in enumerate(zip( 1177 sentence_tokens, sentence_starts, sentence_entities, sentence_events, 1178 sentence_relations, sentence_entity_cluster, sentence_event_cluster, 1179 sentence_relentbase, sentences)): 1180 for entity in ents: 1181 entity.char_offsets_to_token_offsets(toks) 1182 for entity in ent_cls: 1183 for mention in entity.entity_mentions: 1184 mention.char_offsets_to_token_offsets(toks) 1185 for event in evts: 1186 event.trigger.char_offsets_to_token_offsets(toks) 1187 for event in evt_cls: 1188 for mention in event.event_mentions: 1189 mention.trigger.char_offsets_to_token_offsets(toks) 1190 wnd_id = '{}-{}'.format(doc_id, i) 1191 sentence_objs.append(Sentence(start=sent[0][1], 1192 end=sent[-1][2], 1193 text=' '.join([s[0] for s in sent]), 1194 sent_starts=starts, 1195 sent_id=wnd_id, 1196 tokens=[t for _, _, t in toks], 1197 entities=ents, 1198 relations=rels, 1199 events=evts, 1200 entity_cluster=ent_cls, 1201 relations_entbase=rel_ents, 1202 event_cluster=evt_cls)) 1203 return Document(doc_id, sentence_objs) 1204 1205 1206 def convert_batch(input_path: str, 1207 output_path: str, 1208 time_and_val: bool = False, 1209 language: str = 'english', 1210 window_size: int = 1): 1211 """Converts a batch of documents. 1212 1213 Args: 1214 input_path (str): path to the input directory. Usually, it is the path 1215 to the LDC2006T06/data/English or LDC2006T06/data/Chinese folder. 1216 output_path (str): path to the output JSON file. 1217 time_and_val (bool, optional): extracts times and values or not. 1218 Defaults to False. 1219 language (str, optional): document language. Available options: english, 1220 chinese, arabic. Defaults to 'english'. 1221 """ 1222 if language == 'english': 1223 sgm_files = glob.glob(os.path.join( 1224 input_path, '**', 'timex2norm', '*.sgm')) 1225 elif language == 'chinese' or language == 'arabic': 1226 sgm_files = glob.glob(os.path.join( 1227 input_path, '**', 'adj', '*.sgm')) 1228 else: 1229 raise ValueError('Unknown language: {}'.format(language)) 1230 print(input_path) 1231 print('Converting the dataset to JSON format') 1232 print('#SGM files: {}'.format(len(sgm_files))) 1233 progress = tqdm.tqdm(total=len(sgm_files)) 1234 with open(output_path, 'w', encoding='utf-8') as w: 1235 for sgm_file in sgm_files: 1236 progress.update(1) 1237 apf_file = sgm_file.replace('.sgm', '.apf.xml') 1238 doc = convert(sgm_file, apf_file, time_and_val=time_and_val, 1239 language=language, window_size=window_size) 1240 if doc: 1241 w.write(json.dumps(doc.to_dict()) + '\n') 1242 progress.close() 1243 1244 1245 def strQ2B(ustring): 1246 ss = [] 1247 for s in ustring: 1248 rstring = "" 1249 for uchar in s: 1250 inside_code=ord(uchar) 1251 if inside_code==12288: 1252 inside_code=32 1253 elif (65373>= inside_code >= 65281): 1254 inside_code -= 65248 1255 rstring += chr(inside_code) 1256 if rstring == '…': 1257 rstring = '...' 1258 ss.append(rstring) 1259 return ''.join(ss) 1260 1261 1262 def self_define_equal(current, ori_token): 1263 ori_str = strQ2B(ori_token) 1264 cur_str = ''.join([x if x!='<unk>' else '⁔' for x in current]) 1265 if len(ori_str) != len(cur_str): 1266 return False 1267 for o,c in zip(ori_str, cur_str): 1268 if (o != c) and (c != '⁔'): 1269 return False 1270 return True 1271 1272 1273 def map_decode_back_pieces(encoded_input, ori_tokens, tokenizer): 1274 decoded = [tokenizer.decode(x) for x in encoded_input['input_ids']] 1275 pieces = [] 1276 ori_cnt = 0 1277 current = [] 1278 for d in decoded: 1279 if d == '': 1280 continue 1281 current.append(d) 1282 if self_define_equal(current, ori_tokens[ori_cnt]): 1283 pieces.append(current) 1284 current = [] 1285 ori_cnt += 1 1286 assert len(pieces) == len(ori_tokens) 1287 return pieces 1288 1289 1290 def convert_to_oneie(input_path: str, 1291 output_path: str, 1292 language: str, 1293 tokenizer: PreTrainedTokenizer): 1294 """Converts files to OneIE format. 1295 1296 Args: 1297 input_path (str): path to the input file. 1298 output_path (str): path to the output file. 1299 tokenizer (PreTrainedTokenizer): wordpiece tokenizer. 1300 """ 1301 print('Converting the dataset to OneIE format') 1302 skip_num = 0 1303 with open(input_path, 'r', encoding='utf-8') as r, \ 1304 open(output_path, 'w', encoding='utf-8') as w: 1305 for line in r: 1306 doc = json.loads(line) 1307 for sentence in doc['sentences']: 1308 tokens = sentence['tokens'] 1309 if tokenizer.__class__.__name__.startswith('XLMRoberta') and language =='chinese': 1310 encoded_input = tokenizer(tokens, 1311 is_pretokenized=True, 1312 add_special_tokens=False) 1313 pieces = map_decode_back_pieces(encoded_input, tokens, tokenizer) 1314 else: 1315 pieces = [tokenizer.tokenize(t) for t in tokens] 1316 token_lens = [len(x) for x in pieces] 1317 if 0 in token_lens: 1318 skip_num += 1 1319 continue 1320 pieces = [p for ps in pieces for p in ps] 1321 if len(pieces) == 0: 1322 skip_num += 1 1323 continue 1324 1325 entity_text = {e['mention_id']: e['text'] 1326 for e in sentence['entities']} 1327 # update argument text 1328 for relation in sentence['relations']: 1329 arg1, arg2 = relation['arg1'], relation['arg2'] 1330 arg1['text'] = entity_text[arg1['mention_id']] 1331 arg2['text'] = entity_text[arg2['mention_id']] 1332 for event in sentence['events']: 1333 for arg in event['arguments']: 1334 arg['text'] = entity_text[arg['mention_id']] 1335 for event in sentence['event_cluster']: 1336 for mention in event['event_mentions']: 1337 for arg in mention['arguments']: 1338 arg['text'] = entity_text[arg['mention_id']] 1339 1340 # entities 1341 entities = [] 1342 for entity in sentence['entities']: 1343 assert entity['end']-entity['start'] >= 1 1344 entities.append({ 1345 'id': entity['mention_id'], 1346 'text': entity['text'], 1347 'entity_type': entity['entity_type'], 1348 'mention_type': entity['mention_type'], 1349 'entity_subtype': entity['entity_subtype'], 1350 'start': entity['start'], 1351 'end': entity['end'], 1352 }) 1353 1354 # relations 1355 relations = [] 1356 for relation in sentence['relations']: 1357 relations.append({ 1358 'id': relation['relation_id'], 1359 'relation_type': relation['relation_type'], 1360 'relation_subtype': '{}:{}'.format(relation['relation_type'], 1361 relation['relation_subtype']), 1362 'arguments': [ 1363 { 1364 'entity_id': relation['arg1']['mention_id'], 1365 'text': relation['arg1']['text'], 1366 'role': relation['arg1']['role'] 1367 }, 1368 { 1369 'entity_id': relation['arg2']['mention_id'], 1370 'text': relation['arg2']['text'], 1371 'role': relation['arg2']['role'] 1372 } 1373 ] 1374 }) 1375 1376 # events 1377 events = [] 1378 for event in sentence['events']: 1379 events.append({ 1380 'id': event['mention_id'], 1381 'event_type': '{}:{}'.format(event['event_type'], 1382 event['event_subtype']), 1383 'trigger': event['trigger'], 1384 'arguments': [ 1385 { 1386 'entity_id': arg['mention_id'], 1387 'text': arg['text'], 1388 'role': arg['role'] 1389 } for arg in event['arguments'] 1390 ] 1391 }) 1392 1393 # coreference 1394 corefs = [] 1395 for entity in sentence['entity_cluster']: 1396 if len(entity['entity_mentions']) > 1: 1397 corefs.append({ 1398 'id': entity['entity_id'], 1399 'entities': [{ 1400 'id': mention['mention_id'], 1401 'text': mention['text'], 1402 'entity_type': mention['entity_type'], 1403 'mention_type': mention['mention_type'], 1404 'entity_subtype': mention['entity_subtype'], 1405 'start': mention['start'], 1406 'end': mention['end'], 1407 } for mention in entity['entity_mentions']], 1408 'entity_type': entity['entity_type'], 1409 'entity_subtype': entity['entity_subtype'] 1410 }) 1411 1412 event_corefs = [] 1413 for event in sentence['event_cluster']: 1414 if len(event['event_mentions']) > 1: 1415 event_corefs.append({ 1416 'id': event['event_id'], 1417 'events':[{ 1418 'id': mention['mention_id'], 1419 'event_type': '{}:{}'.format(mention['event_type'],mention['event_subtype']), 1420 'trigger': mention['trigger'], 1421 'arguments': [ 1422 { 1423 'entity_id': arg['mention_id'], 1424 'text': arg['text'], 1425 'role': arg['role'] 1426 } for arg in mention['arguments'] 1427 ] 1428 } for mention in event['event_mentions']], 1429 'event_type': '{}:{}'.format(mention['event_type'],mention['event_subtype']), 1430 'event_arguments(entity_base)':[{ 1431 'entity_id': arg['entity_id'], 1432 'role': arg['role'] 1433 } for arg in event['arguments']] 1434 }) 1435 1436 sent_obj = { 1437 'doc_id': doc['doc_id'], 1438 'wnd_id': sentence['sent_id'], 1439 'tokens': tokens, 1440 'pieces': pieces, 1441 'token_lens': token_lens, 1442 'sentence': sentence['text'], 1443 'entity_mentions': entities, 1444 'relation_mentions': relations, 1445 'event_mentions': events, 1446 'entity_coreference': corefs, 1447 'event_coreference': event_corefs, 1448 'sentence_starts': sentence['sentence_starts'][:-1] 1449 } 1450 w.write(json.dumps(sent_obj) + '\n') 1451 print('skip num: {}'.format(skip_num)) 1452 1453 1454 def split_data(input_file: str, 1455 output_dir: str, 1456 split_path: str, 1457 window_size: int): 1458 """Splits the input file into train/dev/test sets. 1459 1460 Args: 1461 input_file (str): path to the input file. 1462 output_dir (str): path to the output directory. 1463 split_path (str): path to the split directory that contains three files, 1464 train.doc.txt, dev.doc.txt, and test.doc.txt . Each line in these 1465 files is a document ID. 1466 """ 1467 print('Splitting the dataset into train/dev/test sets') 1468 train_docs, dev_docs, test_docs = set(), set(), set() 1469 # load doc ids 1470 with open(os.path.join(split_path, 'train.doc.txt')) as r: 1471 train_docs.update(r.read().strip('\n').split('\n')) 1472 with open(os.path.join(split_path, 'dev.doc.txt')) as r: 1473 dev_docs.update(r.read().strip('\n').split('\n')) 1474 with open(os.path.join(split_path, 'test.doc.txt')) as r: 1475 test_docs.update(r.read().strip('\n').split('\n')) 1476 1477 if window_size > 10000: 1478 f_size = 'doc' 1479 else: 1480 f_size = 'w{}'.format(window_size) 1481 1482 # split the dataset 1483 with open(input_file, 'r', encoding='utf-8') as r, \ 1484 open(os.path.join(output_dir, 'train.{}.oneie.json'.format(f_size)), 'w') as w_train, \ 1485 open(os.path.join(output_dir, 'dev.{}.oneie.json'.format(f_size)), 'w') as w_dev, \ 1486 open(os.path.join(output_dir, 'test.{}.oneie.json'.format(f_size)), 'w') as w_test: 1487 for line in r: 1488 inst = json.loads(line) 1489 doc_id = inst['doc_id'] 1490 if doc_id in train_docs: 1491 w_train.write(line) 1492 elif doc_id in dev_docs: 1493 w_dev.write(line) 1494 elif doc_id in test_docs: 1495 w_test.write(line) 1496 else: 1497 print('missing!! {}'.format(doc_id)) 1498 1499 1500 if __name__ == '__main__': 1501 parser = ArgumentParser() 1502 parser.add_argument('-i', '--input', help='Path to the input folder') 1503 parser.add_argument('-o', '--output', help='Path to the output folder') 1504 parser.add_argument('-s', '--split', default=None, 1505 help='Path to the split folder') 1506 parser.add_argument('-b', '--bert', 1507 help='BERT model name', 1508 default='bert-large-cased') 1509 parser.add_argument('-c', '--bert_cache_dir', 1510 help='Path to the BERT cahce directory') 1511 parser.add_argument('-l', '--lang', default='english', 1512 help='Document language') 1513 parser.add_argument('--time_and_val', action='store_true', 1514 help='Extracts times and values') 1515 parser.add_argument('-w', '--window', default=1, help='Integer for window size', type=int) 1516 1517 args = parser.parse_args() 1518 if args.lang not in ['chinese', 'english', 'arabic']: 1519 raise ValueError('Unsupported language: {}'.format(args.lang)) 1520 input_dir = os.path.join(args.input, args.lang.title()) 1521 1522 # Create a tokenizer based on the model name 1523 model_name = args.bert 1524 cache_dir = args.bert_cache_dir 1525 if model_name.startswith('bert-'): 1526 tokenizer = BertTokenizer.from_pretrained(model_name, 1527 cache_dir=cache_dir) 1528 elif model_name.startswith('roberta-'): 1529 tokenizer = RobertaTokenizer.from_pretrained(model_name, 1530 cache_dir=cache_dir) 1531 elif model_name.startswith('xlm-roberta-'): 1532 tokenizer = XLMRobertaTokenizer.from_pretrained(model_name, 1533 cache_dir=cache_dir) 1534 elif model_name.startswith('lanwuwei'): 1535 tokenizer = BertTokenizer.from_pretrained(model_name, 1536 cache_dir=cache_dir, 1537 do_lower_case=True) 1538 else: 1539 tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir, do_lower_case=False, use_fast=False) 1540 #raise ValueError('Unknown model name: {}'.format(model_name)) 1541 1542 if args.window > 10000: 1543 f_size = 'doc' 1544 else: 1545 f_size = 'w{}'.format(args.window) 1546 1547 with torch.no_grad(): 1548 # Convert to doc-level JSON format 1549 json_path = os.path.join(args.output, '{}.{}.json'.format(args.lang, f_size)) 1550 convert_batch(input_dir, json_path, time_and_val=args.time_and_val, 1551 language=args.lang, window_size=args.window) 1552 1553 # Convert to OneIE format 1554 oneie_path = os.path.join(args.output, '{}.{}.oneie.json'.format(args.lang, f_size)) 1555 convert_to_oneie(json_path, oneie_path, args.lang, tokenizer=tokenizer) 1556 1557 # Split the data 1558 if args.split: 1559 split_data(oneie_path, args.output, args.split, args.window)