process_ucdp.py (3740B)
1 import json 2 from argparse import ArgumentParser 3 from transformers import BertTokenizer, RobertaTokenizer, AutoTokenizer 4 5 6 ROLES = [ 7 "side_a_name", 8 "side_b_name", 9 "start_date", 10 "end_date", 11 "location_root_name", 12 "location_adm1_name", 13 "location_adm2_name", 14 "location_where_name", 15 "deaths_side_a", 16 "deaths_side_b", 17 "deaths_civilian", 18 "deaths_unknown", 19 "deaths_low", 20 "deaths_high", 21 ] 22 23 ROLE_TO_TYPE = { 24 "side_a_name": "Actor", 25 "side_b_name": "Actor", 26 "start_date": "Date", 27 "end_date": "Date", 28 "location_root_name": "Place", 29 "location_adm1_name": "Place", 30 "location_adm2_name": "Place", 31 "location_where_name": "Place", 32 "deaths_side_a": "Deaths", 33 "deaths_side_b": "Deaths", 34 "deaths_civilian": "Deaths", 35 "deaths_unknown": "Deaths", 36 "deaths_low": "Deaths", 37 "deaths_high": "Deaths", 38 } 39 40 41 def convert(input_file, output_file, tokenizer): 42 with open(input_file, "r", encoding="utf-8") as r, open(output_file, "w", encoding="utf-8") as w: 43 for line in r: 44 sample = json.loads(line) 45 entities = [ 46 { 47 "id": f"{sample['id']}-{irole}", 48 "start": 2*irole, 49 "end": 2*irole+1, 50 "entity_type": ROLE_TO_TYPE[role], 51 "mention_type": "UNK", 52 "text": str(sample[role]), 53 } 54 for irole, role in enumerate(ROLES) 55 ] 56 event_mentions = [{ 57 "event_type": "UCDPDeath", 58 "id": str(sample["id"]), 59 "trigger": { 60 "start": 0, 61 "end": 1, 62 "text": "UNK", 63 }, 64 "arguments": [ 65 { 66 "entity_id": f"{sample['id']}-{irole}", 67 "role": role, 68 "text": str(sample[role]), 69 } 70 for irole, role in enumerate(ROLES) 71 ] 72 }] 73 input_text = f"{sample['source_date']} {sample['source_article']}" 74 input_tokenized = tokenizer.tokenize(input_text) 75 print(json.dumps({ 76 "doc_id": sample["id"], 77 "wnd_id": 0, 78 "entity_mentions": entities, 79 "relation_mentions": [], 80 "event_mentions": event_mentions, 81 "entity_coreference": [], 82 "event_coreference": [], 83 "pieces": input_tokenized, 84 "tokens": [input_text], 85 "token_lens": [len(input_tokenized)], 86 "sentence": input_text, 87 }), file=w) 88 89 90 if __name__ == '__main__': 91 parser = ArgumentParser() 92 parser.add_argument('-i', '--input', help='Path to the input dataset split') 93 parser.add_argument('-o', '--output', help='Path to the output file') 94 parser.add_argument('-b', '--bert', help='BERT model name', default='bert-large-cased') 95 args = parser.parse_args() 96 model_name = args.bert 97 if model_name.startswith('bert-'): 98 bert_tokenizer = BertTokenizer.from_pretrained(args.bert, 99 do_lower_case=False) 100 elif model_name.startswith('roberta-'): 101 bert_tokenizer = RobertaTokenizer.from_pretrained(args.bert, 102 do_lower_case=False) 103 else: 104 bert_tokenizer = AutoTokenizer.from_pretrained(args.bert, do_lower_case=False, use_fast=False) 105 106 convert(args.input, args.output, bert_tokenizer)