generate_data_degree_e2e.py (4562B)
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 23 # fix random seed 24 random.seed(config.seed) 25 np.random.seed(config.seed) 26 27 # logger 28 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(message)s', datefmt='[%Y-%m-%d %H:%M:%S]') 29 logger = logging.getLogger(__name__) 30 logger.info(f"\n{pprint.pformat(vars(config), indent=4)}") 31 32 def generate_data(data_set, vocab, config): 33 inputs = [] 34 targets = [] 35 events = [] 36 37 for data in data_set.data: 38 event_template = eve_template_generator(data.tokens, data.triggers, data.roles, config.input_style, config.output_style, vocab, True) 39 all_data_ = event_template.get_training_data() 40 pos_data_ = [dt for dt in all_data_ if dt[3]] 41 neg_data_ = [dt for dt in all_data_ if not dt[3]] 42 np.random.shuffle(neg_data_) 43 44 for data_ in pos_data_: 45 inputs.append(data_[0]) 46 targets.append(data_[1]) 47 events.append((data_[2], data_[4], data_[5])) 48 49 for data_ in neg_data_[:config.n_negative]: 50 inputs.append(data_[0]) 51 targets.append(data_[1]) 52 events.append((data_[2], data_[4], data_[5])) 53 54 return inputs, targets, events 55 56 # check valid styles 57 assert np.all([style in ['event_type_sent', 'keywords', 'template'] for style in config.input_style]) 58 assert np.all([style in ['trigger:sentence', 'argument:sentence'] for style in config.output_style]) 59 60 # tokenizer 61 tokenizer = AutoTokenizer.from_pretrained(config.model_name, cache_dir=config.cache_dir) 62 special_tokens = ['<Trigger>', '<sep>'] 63 tokenizer.add_tokens(special_tokens) 64 65 if not os.path.exists(config.finetune_dir): 66 os.makedirs(config.finetune_dir) 67 68 # load data 69 train_set = EEDataset(tokenizer, config.train_file, max_length=config.max_length) 70 dev_set = EEDataset(tokenizer, config.dev_file, max_length=config.max_length) 71 test_set = EEDataset(tokenizer, config.test_file, max_length=config.max_length) 72 vocab = generate_vocabs([train_set, dev_set, test_set]) 73 74 # save vocabulary 75 with open('{}/vocab.json'.format(config.finetune_dir), 'w') as f: 76 json.dump(vocab, f, indent=4) 77 78 # generate finetune data 79 train_inputs, train_targets, train_events = generate_data(train_set, vocab, config) 80 logger.info(f"Generated {len(train_inputs)} training examples from {len(train_set)} instance") 81 82 with open('{}/train_input.json'.format(config.finetune_dir), 'w') as f: 83 json.dump(train_inputs, f, indent=4) 84 85 with open('{}/train_target.json'.format(config.finetune_dir), 'w') as f: 86 json.dump(train_targets, f, indent=4) 87 88 with open('{}/train_all.pkl'.format(config.finetune_dir), 'wb') as f: 89 pickle.dump({ 90 'input': train_inputs, 91 'target': train_targets, 92 'all': train_events 93 }, f) 94 95 dev_inputs, dev_targets, dev_events = generate_data(dev_set, vocab, config) 96 logger.info(f"Generated {len(dev_inputs)} dev examples from {len(dev_set)} instance") 97 98 with open('{}/dev_input.json'.format(config.finetune_dir), 'w') as f: 99 json.dump(dev_inputs, f, indent=4) 100 101 with open('{}/dev_target.json'.format(config.finetune_dir), 'w') as f: 102 json.dump(dev_targets, f, indent=4) 103 104 with open('{}/dev_all.pkl'.format(config.finetune_dir), 'wb') as f: 105 pickle.dump({ 106 'input': dev_inputs, 107 'target': dev_targets, 108 'all': dev_events 109 }, f) 110 111 test_inputs, test_targets, test_events = generate_data(test_set, vocab, config) 112 logger.info(f"Generated {len(test_inputs)} test examples from {len(test_set)} instance") 113 114 with open('{}/test_input.json'.format(config.finetune_dir), 'w') as f: 115 json.dump(test_inputs, f, indent=4) 116 117 with open('{}/test_target.json'.format(config.finetune_dir), 'w') as f: 118 json.dump(test_targets, f, indent=4) 119 120 with open('{}/test_all.pkl'.format(config.finetune_dir), 'wb') as f: 121 pickle.dump({ 122 'input': test_inputs, 123 'target': test_targets, 124 'all': test_events 125 }, f) 126 127 128