evaluate.py

import os
import sys
import time
import copy
import json
import torch
import pickle
import random
import logging
import subprocess
import numpy as np
import jsbeautifier
from args import load_args
import multiprocessing as mp
from logger import create_logger
from collections import OrderedDict
from tensorboardX import SummaryWriter
from models import load_model, save_model
from processors import processor_dict, compute_metrics
from utils import warp_tqdm, compute_confidence_interval
from torch.utils.data import DataLoader, SequentialSampler, RandomSampler, TensorDataset
from few_shot_modules import SequenceClassificationEvaluateLshotBenchmark, SequenceClassificationLshotEvaluate

opts = jsbeautifier.default_options()

def dist_training(args):
    # Setup CUDA, GPU & distributed training
    if args.local_rank == -1 or args.no_cuda:
        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
        args.n_gpu = torch.cuda.device_count()
    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        torch.distributed.init_process_group(backend="nccl")
        args.n_gpu = 1
    args.device = device
    return args

def set_seed(_seed: int):
    random.seed(_seed)
    np.random.seed(_seed)
    torch.manual_seed(_seed)
    torch.cuda.manual_seed_all(_seed)


def SequenceClassificationEvaluate(
        args,
        model, tokenizer, processor,
        file_info,
        conf_penalty=0,
        logger=None,
        mode="dev", 
        verbose=True
    ):
    if logger is None:
        logger = logging.getLogger(__name__)
    if verbose:
        logger.info("Evaluating ...")
    # load data
    eval_dataset = processor.load_and_cache_examples(
                args, tokenizer, mode, file_info, logger
        )
    if verbose:
        logger.info("Data loader creation ...")
    args.eval_batch_size = args.per_gpu_eval_batch_size
    sampler = SequentialSampler(eval_dataset) 
    eval_dataloader = DataLoader(eval_dataset, sampler=sampler, batch_size=args.eval_batch_size)

    if verbose:
        logger.info("***** Running evaluation {} *****".format(file_info))
        logger.info("  Num examples = {}".format( len(eval_dataloader) ))
        logger.info("  Batch size = %d", args.eval_batch_size)
    
    eval_loss, curr_eval_steps, total_number_of_sample = 0.0, 0, 0
    total_number_of_step = len(eval_dataloader)
    lm_outputs, logits, true_label_ids, per_sample_loss = None, None, None, None
    model.eval()
    
    disable_tqdm = True if verbose==True else args.disable_tqdm
    for batch in warp_tqdm(disable_tqdm, file_info, eval_dataloader):
        batch = tuple(t.to(args.device) for t in batch)
        if curr_eval_steps % 20 == 0 and verbose==True:
            logger.info("  Evaluating {}/{}".format(curr_eval_steps, total_number_of_step))
        curr_eval_steps += 1
        with torch.no_grad():
            inputs = {
                        "input_ids": batch[0],
                        "attention_mask": batch[1],
                        "labels": batch[3],
                        "conf_penalty":conf_penalty
                    }
            if args.model_type != "distilbert":
                inputs["token_type_ids"] = batch[2] if args.model_type in ["bert", "xlnet"] else None  # XLM and RoBERTa don"t use segment_ids
            batch_eval_loss, outputs = model(**inputs)
            batch_per_sample_loss = outputs[0]
            batch_logits = outputs[1]
            batch_lm_outputs = outputs[2][:, 0, :]

            if args.n_gpu > 1:
                batch_eval_loss = batch_eval_loss.mean()  # mean() to average on multi-gpu parallel evaluating
            eval_loss += batch_eval_loss.item()

        total_number_of_sample += batch_logits.size()[0]
        if logits is None:
            logits = batch_logits.detach().cpu().numpy()
            true_label_ids = inputs["labels"].detach().cpu().numpy()
            per_sample_loss = list(batch_per_sample_loss.detach().cpu().numpy())
            lm_outputs = batch_lm_outputs.detach().cpu().numpy()
        else:
            logits = np.append(logits, batch_logits.detach().cpu().numpy(), axis=0)
            true_label_ids = np.append(true_label_ids, inputs["labels"].detach().cpu().numpy(), axis=0)
            per_sample_loss = per_sample_loss + list(batch_per_sample_loss.detach().cpu().numpy())
            lm_outputs = np.append(lm_outputs, batch_lm_outputs.detach().cpu().numpy(), axis=0)

    if verbose:
        logger.info("Total number of sample evaluated : {}".format(total_number_of_sample))
    eval_loss = eval_loss / total_number_of_sample
    if args.output_mode == "classification":
        preds = np.argmax(logits, axis=1)
    else:
        raise ValueError("No other `output_mode` for SequenceClassification task.")
    result = compute_metrics(args.task_name, preds, true_label_ids)
    result["loss"] = eval_loss
    if verbose:
        logger.info("***** Eval results :{} *****".format(file_info))
        for key in sorted(result.keys()):
            logger.info("  %s = %s", key, str(result[key]))

    return result, lm_outputs, logits, per_sample_loss



def eval_checkpoint(
        args, 
        model, tokenizer, processor,
        global_step,
        file_infos, eval_func, 
        tb_writer=None,
        tf_board_header="",
        mode="test",
        logger=None
    ):
    if logger is None:
        logger = logging.getLogger(__name__)
    if mode == "dev":
        langs = args.dev_lang 
    elif mode == "test":
        langs = args.tgt_lang 
    elif mode == "train":
        langs = args.src_lang 
    else:
        raise NotImplementedError()

    total_eval_file  = 0
    dataset_results = {}
    if file_infos is not None:
        tot_metric = 0
        for file_info in file_infos:
            if file_info.split(";")[-1] not in langs:
                continue
            total_eval_file += 1
            result, _, _, _ = eval_func(
                            args,
                            model, tokenizer, processor,
                            file_info,
                            logger,
                            mode=mode
                        )
            _ = None # free up memory
            if tb_writer is not None:
                for key, value in result.items():
                    tb_writer.add_scalar(
                        "{}_{}_{}_{}".format(tf_board_header, file_info, mode, key), 
                        value,
                        global_step
                    )
            tot_metric = tot_metric + result[args.model_selection_metric]
            dataset_results[file_info] = result
        avg_metric = tot_metric/float(total_eval_file)
        dataset_results["avg_metric"] = avg_metric
    return dataset_results


def accumulate_report(result_dict, model_selection_metric, logger):
    tot, cnt = 0, 0
    for k, v in result_dict.items():
        metric_val = round(v[0][model_selection_metric]*100, 3)
        tot += metric_val
        cnt += 1
    logger.info("Average : {}".format(round(tot/cnt, 2)))


def baseline_eval(args, logger):
    if logger is None:
        logger = logging.getLogger(__name__)
    if args.task_name not in processor_dict:
        raise ValueError("Task not found in processor pool: {}".format((args.task_name)))
    if args.task_name not in evaluate_func_dict:
        raise ValueError("Task not found eval function pool: {}".format((args.task_name)))
    processor = processor_dict[args.task_name](
                    src_lang=args.src_lang, 
                    dev_lang=args.dev_lang, 
                    tgt_lang=args.tgt_lang, 
                    seed=args.seed,
                    percentage=args.train_data_percentage,
                    n_shot=None,
                    pad_token_label_id=torch.nn.CrossEntropyLoss().ignore_index
                )
    args.output_mode = processor.output_mode

    label_list = processor.get_labels()
    num_labels = len(label_list)
    
    (config, 
     tokenizer, 
     model,
     optimizer, 
     scheduler) = load_model(
                task_name=args.task_name, 
                config_name=args.config_name, 
                tokenizer_name=args.tokenizer_name, 
                model_name_or_path=args.model_name_or_path, 
                num_labels=num_labels, 
                model_type=args.model_type, 
                logger=logger,
                do_lower_case=args.do_lower_case, 
                cache_dir=args.cache_dir,
                is_fp16=args.fp16
        )
    logger.info("Successfully loaded the config : {}, tokenizer : {}, model : {}".format(
                args.config_name, args.tokenizer_name, args.model_name_or_path))

    evaluate_func = evaluate_func_dict[args.task_name]
    model.to(args.device)

    test_result_dict = {}
    for test_file_info in args.test:
        curr_lang = test_file_info.split(";")[-1]
        if curr_lang not in args.tgt_lang.split(";"):
            continue
        test_infer_data = os.path.join(
                            args.output_dir, 
                            "{}.pkl".format(test_file_info.replace("/", "_").replace("\\", "_").replace(";", "_"))
                        )
        test_infer_data_in_model_folder = os.path.join(
                args.model_name_or_path, 
                "{}.pkl".format(test_file_info.replace("/", "_").replace("\\", "_").replace(";", "_"))
            )
        if os.path.isfile(test_infer_data_in_model_folder) and not args.overwrite_cache:
            data_tuple = pickle.load(open(test_infer_data_in_model_folder, "rb"))
            test_result_dict[test_file_info] = data_tuple
            pickle.dump(data_tuple, open(test_infer_data, "wb"))
        if os.path.isfile(test_infer_data) == False or args.overwrite_cache:
            result, lm_outputs_list, logit, per_sample_loss = evaluate_func(
                            args,
                            model, tokenizer, processor,
                            test_file_info,
                            conf_penalty=args.conf_penalty,
                            logger=logger,
                            mode="test", 
                            verbose=False
                        )
            data_tuple = (result, None, logit, per_sample_loss, test_file_info)
            logger.info("{} : {} ({})".format(test_file_info, round(result[args.model_selection_metric]*100, 3), args.model_selection_metric))
            test_result_dict[test_file_info] = data_tuple
            pickle.dump(data_tuple, open(test_infer_data, "wb"))
            pickle.dump(data_tuple, open(test_infer_data_in_model_folder, "wb"))
    accumulate_report(test_result_dict, args.model_selection_metric, logger)
    full_test_res_pkl = os.path.join(args.output_dir, "test_data.pkl")
    pickle.dump(test_result_dict, open(full_test_res_pkl, "wb"))
    full_test_res_pkl_in_model_folder = os.path.join(args.model_name_or_path, "test_data.pkl")
    if not os.path.isfile(full_test_res_pkl_in_model_folder) or not args.overwrite_cache:
        logger.info("Writing ... {}".format(full_test_res_pkl_in_model_folder))
        pickle.dump(test_result_dict, open(full_test_res_pkl_in_model_folder, "wb"))
    

    dev_result_dict = {}
    for dev_file_info in args.dev:
        curr_lang = dev_file_info.split(";")[-1]
        if curr_lang not in args.dev_lang.split(";"):
            continue
        dev_infer_data = os.path.join(
                            args.output_dir, 
                            "{}.pkl".format(dev_file_info.replace("/", "_").replace("\\", "_").replace(";", "_"))
                        )
        dev_infer_data_in_model_folder = os.path.join(
                args.model_name_or_path, 
                "{}.pkl".format(dev_file_info.replace("/", "_").replace("\\", "_").replace(";", "_"))
            )
        if os.path.isfile(dev_infer_data_in_model_folder) and not args.overwrite_cache:
            data_tuple = pickle.load(open(dev_infer_data_in_model_folder, "rb"))
            dev_result_dict[dev_file_info] = data_tuple
            pickle.dump(data_tuple, open(dev_infer_data, "wb"))
        if os.path.isfile(dev_infer_data) == False or args.overwrite_cache:
            result, lm_outputs_list, logit, per_sample_loss = evaluate_func(
                            args,
                            model, tokenizer, processor,
                            dev_file_info,
                            conf_penalty=args.conf_penalty,
                            logger=logger,
                            mode="dev", 
                            verbose=False
                        )
            data_tuple = (result, None, logit, per_sample_loss, dev_file_info)
            logger.info("{} : {} ({})".format(dev_file_info, round(result[args.model_selection_metric]*100, 3), args.model_selection_metric))
            dev_result_dict[dev_file_info] = data_tuple
            pickle.dump(data_tuple, open(dev_infer_data, "wb"))
            pickle.dump(data_tuple, open(dev_infer_data_in_model_folder, "wb"))
    
    accumulate_report(dev_result_dict, args.model_selection_metric, logger)
    full_dev_res_pkl = os.path.join(args.output_dir, "dev_data.pkl")
    pickle.dump(dev_result_dict, open(full_dev_res_pkl, "wb"))
    full_dev_res_pkl_in_model_folder = os.path.join(args.model_name_or_path, "dev_data.pkl")
    if not os.path.isfile(full_dev_res_pkl_in_model_folder) or not args.overwrite_cache:
        logger.info("Writing ... {}".format(full_dev_res_pkl_in_model_folder))
        pickle.dump(dev_result_dict, open(full_dev_res_pkl_in_model_folder, "wb"))

    return 0
        

evaluate_func_dict = {
    "xnli": SequenceClassificationEvaluate,
    "pawsx": SequenceClassificationEvaluate,
    "xnli_few_shot_benchmark": SequenceClassificationEvaluateLshotBenchmark,
    "pawsx_few_shot_benchmark": SequenceClassificationEvaluateLshotBenchmark
}

def load_source_trained_model(args, processor, logger):
    label_list = processor.get_labels()
    num_labels = len(label_list) if args.overwrite_num_of_label is None else args.overwrite_num_of_label
    (config, 
     tokenizer, 
     model,
     optimizer, 
     scheduler) = load_model(
                task_name=args.task_name if args.cross_task_name is None else args.cross_task_name, 
                config_name=args.config_name, 
                tokenizer_name=args.tokenizer_name, 
                model_name_or_path=args.model_name_or_path, 
                num_labels=num_labels, 
                model_type=args.model_type, 
                logger=logger,
                do_lower_case=args.do_lower_case, 
                cache_dir=args.cache_dir,
                is_fp16=args.fp16, 
                verbose=False,
                secondary_num_labels=None if args.overwrite_num_of_label is None else len(label_list)
        )
    logger.info("Successfully loaded the config : {}, tokenizer : {}, model : {}".format(
                args.config_name, args.tokenizer_name, args.model_name_or_path))
    return model, tokenizer


def FewShotBenchmark(args, logger):
    if args.support_split == 'dev':
        args.support_set = args.dev
    elif args.support_split == 'train':
        args.support_set = args.train
    elif args.support_split == 'test':
        args.support_set = args.test
    if logger is None:
        logger = logging.getLogger(__name__)
    if args.task_name not in processor_dict:
        raise ValueError("Task not found in processor pool: {}".format((args.task_name)))
    if args.task_name not in evaluate_func_dict:
        raise ValueError("Task not found eval function pool: {}".format((args.task_name)))
    processor = processor_dict[args.task_name](
                    src_lang=args.src_lang, 
                    dev_lang=args.dev_lang, 
                    tgt_lang=args.tgt_lang, 
                    seed=args.seed,
                    percentage=args.train_data_percentage,
                    n_shot=None,
                    pad_token_label_id=torch.nn.CrossEntropyLoss().ignore_index
                )
    args.output_mode = processor.output_mode
    
    evaluate_func = evaluate_func_dict[args.task_name+"_few_shot_benchmark"]
    base_model, tokenizer = load_source_trained_model(args, processor, logger)

    test_result_dict = {}
    tot_lang = 0
    for test_file_info in args.test:
        curr_lang = test_file_info.split(";")[-1]
        if curr_lang not in args.tgt_lang.split(";"):
            continue
        support_file_info = None
        for file_info in args.support_set:
            support_lang = file_info.split(";")[-1]
            if support_lang not in args.tgt_lang.split(";"):
                continue
            if support_lang == curr_lang:
                support_file_info = file_info
                break
        assert support_file_info is not None
        benchmark_result_file = os.path.join(
                            args.output_dir, 
                            "{}.benchmark.pkl".format(test_file_info.replace("/", "_").replace("\\", "_").replace(";", "_"))
                        )
        result_dict = evaluate_func(
                            args,
                            base_model, 
                            tokenizer,
                            processor,
                            support_file_info,
                            test_file_info,
                            conf_penalty=args.conf_penalty,
                            logger=logger,
                            verbose=True
                        )

        pickle.dump(result_dict, open(benchmark_result_file, "wb"))
        

if __name__ == "__main__":
    main()