#!/usr/bin/env python3 import os from text.torchtext.datasets.generic import Query from text import torchtext from argparse import ArgumentParser import ujson as json import torch import numpy as np import random from pprint import pformat from util import get_splits, set_seed, preprocess_examples from metrics import compute_metrics import models def get_all_splits(args, new_vocab): splits = [] for task in args.tasks: print(f'Loading {task}') kwargs = {} if not 'train' in args.evaluate: kwargs['train'] = None if not 'valid' in args.evaluate: kwargs['validation'] = None if not 'test' in args.evaluate: kwargs['test'] = None s = get_splits(args, task, new_vocab, **kwargs)[0] preprocess_examples(args, [task], [s], new_vocab, train=False) splits.append(s) return splits def prepare_data(args, FIELD): new_vocab = torchtext.data.ReversibleField(batch_first=True, init_token='', eos_token='', lower=args.lower, include_lengths=True) splits = get_all_splits(args, new_vocab) new_vocab.build_vocab(*splits) print(f'Vocabulary has {len(FIELD.vocab)} tokens from training') args.max_generative_vocab = min(len(FIELD.vocab), args.max_generative_vocab) FIELD.append_vocab(new_vocab) print(f'Vocabulary has expanded to {len(FIELD.vocab)} tokens') char_vectors = torchtext.vocab.CharNGram(cache=args.embeddings) glove_vectors = torchtext.vocab.GloVe(cache=args.embeddings) vectors = [char_vectors, glove_vectors] FIELD.vocab.load_vectors(vectors, True) FIELD.decoder_to_vocab = {idx: FIELD.vocab.stoi[word] for idx, word in enumerate(FIELD.decoder_itos)} FIELD.vocab_to_decoder = {idx: FIELD.decoder_stoi[word] for idx, word in enumerate(FIELD.vocab.itos) if word in FIELD.decoder_stoi} splits = get_all_splits(args, FIELD) return FIELD, splits def to_iter(data, bs, device): Iterator = torchtext.data.Iterator it = Iterator(data, batch_size=bs, device=device, batch_size_fn=None, train=False, repeat=False, sort=None, shuffle=None, reverse=False) return it def run(args, field, val_sets, model): device = set_seed(args) print(f'Preparing iterators') if len(args.val_batch_size) == 1 and len(val_sets) > 1: args.val_batch_size *= len(val_sets) iters = [(name, to_iter(x, bs, device)) for name, x, bs in zip(args.tasks, val_sets, args.val_batch_size)] def mult(ps): r = 0 for p in ps: this_r = 1 for s in p.size(): this_r *= s r += this_r return r params = list(filter(lambda p: p.requires_grad, model.parameters())) num_param = mult(params) print(f'{args.model} has {num_param:,} parameters') model.to(device) model.eval() with torch.no_grad(): for task, it in iters: print(task) prediction_file_name = os.path.join(os.path.splitext(args.best_checkpoint)[0], args.evaluate, task + '.txt') answer_file_name = os.path.join(os.path.splitext(args.best_checkpoint)[0], args.evaluate, task + '.gold.txt') results_file_name = answer_file_name.replace('gold', 'results') if 'sql' in task or 'squad' in task: ids_file_name = answer_file_name.replace('gold', 'ids') if os.path.exists(prediction_file_name): print('** ', prediction_file_name, ' already exists -- this is where predictions are stored **') if os.path.exists(answer_file_name): print('** ', answer_file_name, ' already exists -- this is where ground truth answers are stored **') if os.path.exists(results_file_name): print('** ', results_file_name, ' already exists -- this is where metrics are stored **') with open(results_file_name) as results_file: for l in results_file: print(l) if not 'schema' in task and not args.overwrite_predictions and args.silent: continue for x in [prediction_file_name, answer_file_name, results_file_name]: os.makedirs(os.path.dirname(x), exist_ok=True) if not os.path.exists(prediction_file_name) or args.overwrite_predictions: with open(prediction_file_name, 'w') as prediction_file: predictions = [] ids = [] for batch_idx, batch in enumerate(it): _, p = model(batch) p = field.reverse(p) for i, pp in enumerate(p): if 'sql' in task: ids.append(int(batch.wikisql_id[i])) if 'squad' in task: ids.append(it.dataset.q_ids[int(batch.squad_id[i])]) prediction_file.write(pp + '\n') predictions.append(pp) else: with open(prediction_file_name) as prediction_file: predictions = [x.strip() for x in prediction_file.readlines()] if 'sql' in task: with open(ids_file_name, 'w') as id_file: for i in ids: id_file.write(json.dumps(i) + '\n') if 'squad' in task: with open(ids_file_name, 'w') as id_file: for i in ids: id_file.write(i + '\n') def from_all_answers(an): return [it.dataset.all_answers[sid] for sid in an.tolist()] if not os.path.exists(answer_file_name): with open(answer_file_name, 'w') as answer_file: answers = [] for batch_idx, batch in enumerate(it): if hasattr(batch, 'wikisql_id'): a = from_all_answers(batch.wikisql_id.data.cpu()) elif hasattr(batch, 'squad_id'): a = from_all_answers(batch.squad_id.data.cpu()) elif hasattr(batch, 'woz_id'): a = from_all_answers(batch.woz_id.data.cpu()) else: a = field.reverse(batch.answer.data) for aa in a: answers.append(aa) answer_file.write(json.dumps(aa) + '\n') else: with open(answer_file_name) as answer_file: answers = [json.loads(x.strip()) for x in answer_file.readlines()] if len(answers) > 0: if not os.path.exists(results_file_name): metrics, answers = compute_metrics(predictions, answers, bleu='iwslt' in task or 'multi30k' in task or args.bleu, dialogue='woz' in task, rouge='cnn' in task or 'dailymail' in task or args.rouge, logical_form='sql' in task, corpus_f1='zre' in task, args=args) with open(results_file_name, 'w') as results_file: results_file.write(json.dumps(metrics) + '\n') else: with open(results_file_name) as results_file: metrics = json.loads(results_file.readlines()[0]) if not args.silent: for i, (p, a) in enumerate(zip(predictions, answers)): print(f'Prediction {i+1}: {p}\nAnswer {i+1}: {a}\n') print(metrics) def get_args(): parser = ArgumentParser() parser.add_argument('--path', required=True) parser.add_argument('--evaluate', type=str, required=True) parser.add_argument('--tasks', default=['squad', 'iwslt.en.de', 'cnn_dailymail', 'multinli.in.out', 'sst', 'srl', 'zre', 'woz.en', 'wikisql', 'schema'], nargs='+') parser.add_argument('--devices', default=[0], nargs='+', type=int, help='a list of devices that can be used (multi-gpu currently WIP)') parser.add_argument('--seed', default=123, type=int, help='Random seed.') parser.add_argument('--data', default='/decaNLP/.data/', type=str, help='where to load data from.') parser.add_argument('--embeddings', default='/decaNLP/.embeddings', type=str, help='where to save embeddings.') parser.add_argument('--checkpoint_name') parser.add_argument('--bleu', action='store_true', help='whether to use the bleu metric (always on for iwslt)') parser.add_argument('--rouge', action='store_true', help='whether to use the bleu metric (always on for cnn, dailymail, and cnn_dailymail)') parser.add_argument('--overwrite_predictions', action='store_true', help='whether to overwrite previously written predictions') parser.add_argument('--silent', action='store_true', help='whether to print predictions to stdout') args = parser.parse_args() with open(os.path.join(args.path, 'config.json')) as config_file: config = json.load(config_file) retrieve = ['model', 'transformer_layers', 'rnn_layers', 'transformer_hidden', 'dimension', 'load', 'max_val_context_length', 'val_batch_size', 'transformer_heads', 'max_output_length', 'max_generative_vocab', 'lower', 'cove', 'intermediate_cove'] for r in retrieve: if r in config: setattr(args, r, config[r]) elif 'cove' in r: setattr(args, r, False) else: setattr(args, r, None) args.dropout_ratio = 0.0 args.task_to_metric = {'cnn_dailymail': 'avg_rouge', 'iwslt.en.de': 'bleu', 'multinli.in.out': 'em', 'squad': 'nf1', 'srl': 'nf1', 'sst': 'em', 'wikisql': 'lfem', 'woz.en': 'joint_goal_em', 'zre': 'corpus_f1', 'schema': 'em'} if os.path.exists(os.path.join(args.path, 'process_0.log')): args.best_checkpoint = get_best(args) else: args.best_checkpoint = os.path.join(args.path, args.checkpoint_name) return args def get_best(args): with open(os.path.join(args.path, 'config.json')) as f: save_every = json.load(f)['save_every'] with open(os.path.join(args.path, 'process_0.log')) as f: lines = f.readlines() best_score = 0 best_it = 0 deca_scores = {} for l in lines: if 'val' in l: try: task = l.split('val_')[1].split(':')[0] except Exception as e: print(e) continue it = int(l.split('iteration_')[1].split(':')[0]) metric = args.task_to_metric[task] score = float(l.split(metric+'_')[1].split(':')[0]) if it in deca_scores: deca_scores[it]['deca'] += score deca_scores[it][metric] = score else: deca_scores[it] = {'deca': score, metric: score} if deca_scores[it]['deca'] > best_score: best_score = deca_scores[it]['deca'] best_it = it print(best_it) print(best_score) return os.path.join(args.path, f'iteration_{int(best_it)}.pth') if __name__ == '__main__': args = get_args() print(f'Arguments:\n{pformat(vars(args))}') np.random.seed(args.seed) random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) print(f'Loading from {args.best_checkpoint}') save_dict = torch.load(args.best_checkpoint) field = save_dict['field'] print(f'Initializing Model') Model = getattr(models, args.model) model = Model(field, args) model_dict = save_dict['model_state_dict'] backwards_compatible_cove_dict = {} for k, v in model_dict.items(): if 'cove.rnn.' in k: k = k.replace('cove.rnn.', 'cove.rnn1.') backwards_compatible_cove_dict[k] = v model_dict = backwards_compatible_cove_dict model.load_state_dict(model_dict) field, splits = prepare_data(args, field) model.set_embeddings(field.vocab.vectors) run(args, field, splits, model)