predict.py

#!/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='<init>', eos_token='<eos>', 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)

    decaScore = []
    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 args.overwrite:
                    print('**** overwriting ', prediction_file_name, ' ****')
            if os.path.exists(answer_file_name):
                print('** ', answer_file_name, ' already exists -- this is where ground truth answers are stored **')
                if args.overwrite:
                    print('**** overwriting ', answer_file_name, ' ****')
            if os.path.exists(results_file_name):
                print('** ', results_file_name, ' already exists -- this is where metrics are stored **')
                if args.overwrite:
                    print('**** overwriting ', results_file_name, ' ****')
                else:
                    with open(results_file_name) as results_file:
                        if not args.silent:
                            for l in results_file:
                                print(l)
                        metrics = json.loads(results_file.readlines()[0])
                        decaScore.append(metrics[args.task_to_metric[task]])
                    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:
                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) 
                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')
            else:
                with open(prediction_file_name) as prediction_file:
                    predictions = [x.strip() for x in prediction_file.readlines()] 
                if 'sql' in task or 'squad' in task:
                    with open(ids_file_name) as id_file:
                        ids = [int(x.strip()) for x in id_file.readlines()]
   
            def from_all_answers(an):
                return [it.dataset.all_answers[sid] for sid in an.tolist()] 
    
            if not os.path.exists(answer_file_name) or args.overwrite:
                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) or args.overwrite:
                    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)
                decaScore.append(metrics[args.task_to_metric[task]])

    print(f'Evaluated Tasks:\n')
    for i, (task, _) in enumerate(iters):
        print(f'{task}: {decaScore[i]}')
    print(f'-------------------')
    print(f'DecaScore:  {sum(decaScore)}\n')
    print(f'\nSummary: | {sum(decaScore)} | {" | ".join([str(x) for x in decaScore])} |\n')


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', 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', 'elmo', 'glove_and_char']
        for r in retrieve:
            if r in config:
                setattr(args, r,  config[r])
            elif 'cove' in r:
                setattr(args, r, False)
            elif 'elmo' in r:
                setattr(args, r, [-1])
            elif 'glove_and_char' in r:
                setattr(args, r, True)
            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 not args.checkpoint_name is None:
        args.best_checkpoint = os.path.join(args.path, args.checkpoint_name)
    else:
        assert os.path.exists(os.path.join(args.path, 'process_0.log'))
        args.best_checkpoint = get_best(args)
           
    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)