generate_data_degree_eae.py (4234B)
1 import os, json, pickle, logging, pprint, random 2 import numpy as np 3 from data import EEDataset 4 from argparse import ArgumentParser, Namespace 5 from utils import generate_vocabs 6 from transformers import AutoTokenizer 7 import ipdb 8 9 # configuration 10 parser = ArgumentParser() 11 parser.add_argument('-c', '--config', required=True) 12 args = parser.parse_args() 13 with open(args.config) as fp: 14 config = json.load(fp) 15 config.update(args.__dict__) 16 config = Namespace(**config) 17 18 if config.dataset == "ace05e" or config.dataset == "ace05ep": 19 from template_generate_ace import eve_template_generator 20 elif config.dataset == "ere": 21 from template_generate_ere import eve_template_generator 22 elif config.dataset == "ucdp": 23 from template_generate_ucdp import eve_template_generator 24 25 # fix random seed 26 random.seed(config.seed) 27 np.random.seed(config.seed) 28 29 # logger 30 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(message)s', datefmt='[%Y-%m-%d %H:%M:%S]') 31 logger = logging.getLogger(__name__) 32 logger.info(f"\n{pprint.pformat(vars(config), indent=4)}") 33 34 def generate_data(data_set, vocab, config): 35 inputs = [] 36 targets = [] 37 events = [] 38 39 for data in data_set.data: 40 event_template = eve_template_generator(data.tokens, data.triggers, data.roles, config.input_style, config.output_style, vocab, False) 41 for data_ in event_template.get_training_data(): 42 inputs.append(data_[0]) 43 targets.append(data_[1]) 44 events.append(data_[2]) 45 46 return inputs, targets, events 47 48 # check valid styles 49 assert np.all([style in ['event_type_sent', 'triggers', 'template'] for style in config.input_style]) 50 assert np.all([style in ['argument:sentence'] for style in config.output_style]) 51 52 # tokenizer 53 tokenizer = AutoTokenizer.from_pretrained(config.model_name, cache_dir=config.cache_dir) 54 special_tokens = ['<Trigger>', '<sep>'] 55 tokenizer.add_tokens(special_tokens) 56 57 if not os.path.exists(config.finetune_dir): 58 os.makedirs(config.finetune_dir) 59 60 # load data 61 train_set = EEDataset(tokenizer, config.train_file, max_length=config.max_length) 62 dev_set = EEDataset(tokenizer, config.dev_file, max_length=config.max_length) 63 test_set = EEDataset(tokenizer, config.test_file, max_length=config.max_length) 64 vocab = generate_vocabs([train_set, dev_set, test_set]) 65 66 # save vocabulary 67 with open('{}/vocab.json'.format(config.finetune_dir), 'w') as f: 68 json.dump(vocab, f, indent=4) 69 70 # generate finetune data 71 train_inputs, train_targets, train_events = generate_data(train_set, vocab, config) 72 logger.info(f"Generated {len(train_inputs)} training examples from {len(train_set)} instance") 73 74 with open('{}/train_input.json'.format(config.finetune_dir), 'w') as f: 75 json.dump(train_inputs, f, indent=4) 76 77 with open('{}/train_target.json'.format(config.finetune_dir), 'w') as f: 78 json.dump(train_targets, f, indent=4) 79 80 with open('{}/train_all.pkl'.format(config.finetune_dir), 'wb') as f: 81 pickle.dump({ 82 'input': train_inputs, 83 'target': train_targets, 84 'all': train_events 85 }, f) 86 87 dev_inputs, dev_targets, dev_events = generate_data(dev_set, vocab, config) 88 logger.info(f"Generated {len(dev_inputs)} dev examples from {len(dev_set)} instance") 89 90 with open('{}/dev_input.json'.format(config.finetune_dir), 'w') as f: 91 json.dump(dev_inputs, f, indent=4) 92 93 with open('{}/dev_target.json'.format(config.finetune_dir), 'w') as f: 94 json.dump(dev_targets, f, indent=4) 95 96 with open('{}/dev_all.pkl'.format(config.finetune_dir), 'wb') as f: 97 pickle.dump({ 98 'input': dev_inputs, 99 'target': dev_targets, 100 'all': dev_events 101 }, f) 102 103 test_inputs, test_targets, test_events = generate_data(test_set, vocab, config) 104 logger.info(f"Generated {len(test_inputs)} test examples from {len(test_set)} instance") 105 106 with open('{}/test_input.json'.format(config.finetune_dir), 'w') as f: 107 json.dump(test_inputs, f, indent=4) 108 109 with open('{}/test_target.json'.format(config.finetune_dir), 'w') as f: 110 json.dump(test_targets, f, indent=4) 111 112 with open('{}/test_all.pkl'.format(config.finetune_dir), 'wb') as f: 113 pickle.dump({ 114 'input': test_inputs, 115 'target': test_targets, 116 'all': test_events 117 }, f) 118 119 120