few_shot_modules.py

import os
import time
import copy
import tqdm
import torch
import pickle
import random
import logging
import collections
import numpy as np
from scipy import sparse
from scipy.stats import mode
from models import load_model
from numpy import linalg as LA
import torch.nn.functional as F
from matplotlib import pyplot as plt
from processors import compute_metrics
from tensorboardX import SummaryWriter
from sklearn.metrics import accuracy_score
from sklearn.neighbors import NearestNeighbors
from scipy.optimize import linear_sum_assignment
from transformers import get_constant_schedule, AdamW
from utils import warp_tqdm, compute_confidence_interval
from torch.utils.data import DataLoader, SequentialSampler, RandomSampler, TensorDataset

from transformers import (
    WEIGHTS_NAME,
    AdamW,
    get_linear_schedule_with_warmup
)


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

def BULK_seqClassify_few_shot_extract_feature(args, model, src_ref_dataset, tgt_ref_dataset, tgt_test_dataset, src_lang, tgt_lang, force_load = False, prefix=''):
    model_hash = args.model_name_or_path.replace("/","_").replace("\\", "_")
    save_dir = os.path.join(args.output_dir, model_hash)
    os.makedirs(save_dir, exist_ok=True)
    model.eval()
    args.eval_batch_size = args.per_gpu_eval_batch_size
    with torch.no_grad():
        # get reference training mean
        if not os.path.isfile( os.path.join(save_dir,'{}_{}.src_ref_dt_mean.plk'.format(prefix, src_lang)) ) or (force_load == True):
            out_mean, fc_out_mean = [], []
            sampler = SequentialSampler(src_ref_dataset) 
            src_ref_dataloader = DataLoader(src_ref_dataset, sampler=sampler, batch_size=args.eval_batch_size)
            for batch in warp_tqdm(False, src_lang, src_ref_dataloader):
                batch = tuple(t.to(args.device) for t in batch)
                with torch.no_grad():
                    inputs = {
                                "input_ids": batch[0],
                                "attention_mask": batch[1],
                                "sec_classifier": None if args.overwrite_num_of_label is None else True
                            }
                    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)
                    fc_outputs = outputs[1]
                    outputs = outputs[2][:, 0, :]
#                     print(outputs.shape)
                    out_mean.append(outputs.cpu().data.numpy())
                    fc_out_mean.append(fc_outputs.cpu().data.numpy())
#                     if len(out_mean) == 50:
#                         temp_out_mean = np.concatenate(out_mean, axis=0).mean(0)
#                         temp_fc_out_mean = np.concatenate(fc_out_mean, axis=0).mean(0)
#                         out_mean = [temp_out_mean]
#                         fc_out_mean = [temp_fc_out_mean]
                    
            out_mean = np.concatenate(out_mean, axis=0).mean(0)
            fc_out_mean = np.concatenate(fc_out_mean, axis=0).mean(0)
            pickle.dump(
                    [out_mean, fc_out_mean], 
                    open(os.path.join(save_dir,'{}_{}.src_ref_dt_mean.plk'.format(prefix, src_lang)), "wb")
            )
        else:
            out_mean, fc_out_mean = pickle.load(open(os.path.join(save_dir,'{}_{}.src_ref_dt_mean.plk'.format(prefix, src_lang)), "rb"))
        src_ref_out_mean = out_mean
        src_ref_fc_out_mean = fc_out_mean
        
        # get reference target dictionary (LM representation and logit).
        if not os.path.isfile( os.path.join(save_dir,'{}_{}.tgt_ref_dt_dict.plk'.format(prefix, tgt_lang)) ) or (force_load == True):
            output_dict = collections.defaultdict(list)
            fc_output_dict = collections.defaultdict(list)
            sampler = SequentialSampler(tgt_ref_dataset) 
            tgt_ref_dataloader = DataLoader(tgt_ref_dataset, sampler=sampler, batch_size=args.eval_batch_size)
            for batch in warp_tqdm(False, tgt_lang, tgt_ref_dataloader):
                batch = tuple(t.to(args.device) for t in batch)
                with torch.no_grad():
                    inputs = {
                                "input_ids": batch[0],
                                "attention_mask": batch[1],
                                "sec_classifier": None if args.overwrite_num_of_label is None else True
                            }
                    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)
                    fc_outputs = outputs[1]
                    outputs = outputs[2][:, 0, :]
                    outputs = outputs.cpu().data.numpy()
                    fc_outputs = fc_outputs.cpu().data.numpy()
                    for out, fc_out, label in zip(outputs, fc_outputs, batch[3]):
                        output_dict[label.item()].append(out)
                        fc_output_dict[label.item()].append(fc_out)
            pickle.dump(
                    [output_dict, fc_output_dict], 
                    open(os.path.join(save_dir,'{}_{}.tgt_ref_dt_dict.plk'.format(prefix, tgt_lang)), "wb")
            )
        else:
            output_dict, fc_output_dict = pickle.load(open(os.path.join(save_dir, '{}_{}.tgt_ref_dt_dict.plk'.format(prefix, tgt_lang)), "rb"))
        tgt_ref_out_dict = output_dict
        tgt_ref_fc_out_dict = fc_output_dict
        output_dict = None
        fc_output_dict = None
        
        # get target test dictionary (LM representation and logit).
        if not os.path.isfile( os.path.join(save_dir,'{}_{}.tgt_test_dt_dict.plk'.format(prefix, tgt_lang)) ) or (force_load == True):
            output_dict = collections.defaultdict(list)
            fc_output_dict = collections.defaultdict(list)
            sampler = SequentialSampler(tgt_test_dataset) 
            tgt_test_dataloader = DataLoader(tgt_test_dataset, sampler=sampler, batch_size=args.eval_batch_size)
            for batch in warp_tqdm(False, tgt_lang, tgt_test_dataloader):
                batch = tuple(t.to(args.device) for t in batch)
                with torch.no_grad():
                    inputs = {
                                "input_ids": batch[0],
                                "attention_mask": batch[1],
                                "sec_classifier": None if args.overwrite_num_of_label is None else True
                    }
                    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)
                    fc_outputs = outputs[1]
                    outputs = outputs[2][:, 0, :]
                    outputs = outputs.cpu().data.numpy()
                    fc_outputs = fc_outputs.cpu().data.numpy()
                    for out, fc_out, label in zip(outputs, fc_outputs, batch[3]):
                        output_dict[label.item()].append(out)
                        fc_output_dict[label.item()].append(fc_out)
            pickle.dump(
                    [output_dict, fc_output_dict], 
                    open(os.path.join(save_dir,'{}_{}.tgt_test_dt_dict.plk'.format(prefix, tgt_lang)), "wb")
            )
        else:
            output_dict, fc_output_dict = pickle.load(open(os.path.join(save_dir, '{}_{}.tgt_test_dt_dict.plk'.format(prefix, tgt_lang)), "rb"))
        tgt_test_out_dict = output_dict
        tgt_test_fc_out_dict = fc_output_dict
    return (src_ref_out_mean, src_ref_fc_out_mean,
            tgt_ref_out_dict, tgt_ref_fc_out_dict,
            tgt_test_out_dict, tgt_test_fc_out_dict)

def SequenceClassificationLshotEvaluate():
    raise NotImplementedError

def get_optimizer(model, weight_decay, learning_rate, adam_epsilon, optimization, is_sec_classifier):
    no_decay = ["bias", "LayerNorm.weight"]
    if optimization == "full":
        optimizer_grouped_parameters = [
            {
                "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
                "weight_decay": weight_decay
            },
            {
                "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 
                "weight_decay": 0.0
            }
        ]
    elif optimization == "head":
        if is_sec_classifier:
            head_param_name = ["sec_classifier"]
        else:
            head_param_name = ["classifier"]
        params, param_names = [], []
        for n, p in model.named_parameters():
            flag = 0
            for nd in head_param_name:
                if nd in n:
                    flag = 1       
            flag2 = 0
            for nd in no_decay:
                if nd not in n:
                    flag2 += 1 
            if flag == 1 and flag2 == len(no_decay):
                params.append(p)
                param_names.append(n)
        no_decay_params = [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)]
        # no_decay_params_name = [n for n, p in model.named_parameters() if any(nd in n for nd in no_decay)]
        optimizer_grouped_parameters = [
            {
                "params": params,
                "weight_decay": weight_decay
            },
            {
                "params": no_decay_params, 
                "weight_decay": 0.0
            }
        ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate, eps=adam_epsilon)
    return optimizer


def meta_sample_data(
        tgt_test_out_dict, support_out_dict, 
        meta_val_way, meta_val_query, shot
    ):
    support_sample_class = sorted(random.sample(list(support_out_dict.keys()), meta_val_way))
    train_input, train_label = [], []
    for each_class in support_sample_class: 
        total_samples = len(support_out_dict[each_class])
        sample_list = random.sample(list(support_out_dict[each_class]) , total_samples)
        samples = []
        for a in sample_list:
            samples += [a]
        train_label += [each_class] * len(samples)
        train_input += samples
        
    test_sample_class = support_sample_class
    test_input, test_label = [], []
    for each_class in test_sample_class: 
        total_samples = len(tgt_test_out_dict[each_class])
        sample_list = random.sample(list(tgt_test_out_dict[each_class]), total_samples)
        samples = []
        for a in sample_list:
            samples += [a]
        test_label += [each_class] * len(samples)
        test_input += samples

    train_input = np.array(train_input).astype(np.float32)
    test_input = np.array(test_input).astype(np.float32)
    train_label = np.array(train_label)
    test_label = np.array(test_label)
    return train_input, train_label, test_input, test_label


def create_affinity(X, knn):
    N, D = X.shape
    # print('Compute Affinity ')
    nbrs = NearestNeighbors(n_neighbors=knn).fit(X)
    dist, knnind = nbrs.kneighbors(X)

    row = np.repeat(range(N), knn - 1)
    col = knnind[:, 1:].flatten()
    data = np.ones(X.shape[0] * (knn - 1))
    W = sparse.csc_matrix((data, (row, col)), shape=(N, N), dtype=np.float)
    return W


def normalize(Y_in):
    maxcol = np.max(Y_in, axis=1)
    Y_in = Y_in - maxcol[:, np.newaxis]
    N = Y_in.shape[0]
    size_limit = 150000
    if N > size_limit:
        batch_size = 1280
        Y_out = []
        num_batch = int(math.ceil(1.0 * N / batch_size))
        for batch_idx in range(num_batch):
            start = batch_idx * batch_size
            end = min((batch_idx + 1) * batch_size, N)
            tmp = np.exp(Y_in[start:end, :])
            tmp = tmp / (np.sum(tmp, axis=1)[:, None])
            Y_out.append(tmp)
        del Y_in
        Y_out = np.vstack(Y_out)
    else:
        Y_out = np.exp(Y_in)
        Y_out = Y_out / (np.sum(Y_out, axis=1)[:, None])

    return Y_out


def entropy_energy(Y, unary, kernel, bound_lambda, batch=False):
    tot_size = Y.shape[0]
    pairwise = kernel.dot(Y)
    if batch == False:
        temp = (unary * Y) + (-bound_lambda * pairwise * Y)
        E = (Y * np.log(np.maximum(Y, 1e-20)) + temp).sum()
    else:
        batch_size = 1024
        num_batch = int(math.ceil(1.0 * tot_size / batch_size))
        E = 0
        for batch_idx in range(num_batch):
            start = batch_idx * batch_size
            end = min((batch_idx + 1) * batch_size, tot_size)
            temp = (unary[start:end] * Y[start:end]) + (-bound_lambda * pairwise[start:end] * Y[start:end])
            E = E + (Y[start:end] * np.log(np.maximum(Y[start:end], 1e-20)) + temp).sum()

    return E


def plot_convergence(filename, E_list):
    # Plot the convergence of bound updates

    E_list = np.asarray(E_list)
    length = len(E_list)
    iter_range = np.asarray(list(range(length)))
    plt.figure(1, figsize=(6.4, 4.8))
    plt.ion()
    plt.clf()
    ylabel = r'$\mathcal{B}_i(\mathbf{Y})$'
    plt.plot(iter_range, E_list, 'b-', linewidth=2.2)
    plt.xticks(iter_range[1::2],(iter_range[1::2]+1))
    plt.xlabel('iterations')
    plt.ylabel(ylabel)
    plt.savefig(filename, format='png', dpi=800, bbox_inches='tight')
    plt.show()
    plt.close('all')


def bound_update(args, unary, kernel, bound_lambda, bound_iteration=20, batch=False):
    """
    """
    oldE = float('inf')
    Y = normalize(-unary)
    if float(bound_lambda) == 0.0:
        l = np.argmax(Y, axis=1)
        return l, Y

    E_list = []
    for i in range(bound_iteration):
        additive = -unary
        mul_kernel = kernel.dot(Y)
        Y = -bound_lambda * mul_kernel
        additive = additive - Y
        Y = normalize(additive)
        E = entropy_energy(Y, unary, kernel, bound_lambda, batch)
        E_list.append(E)
        # print('entropy_energy is ' +repr(E) + ' at iteration ',i)
        if (i > 1 and (abs(E - oldE) <= 1e-6 * abs(oldE))):
            # print('Converged')
            break

        else:
            oldE = E.copy()

    if args.plot_converge:
      filename = os.path.join(args.output_dir,'convergence_{}.png'.format(args.task_name))
      plot_convergence(filename,E_list)

    l = np.argmax(Y, axis=1)
    return l, Y


def get_accuracy(L1, L2):
    # Since the labels may be different we utilize the Hungarian method to ensure the map of
    # the original ground truth labeling with the returned labels from our laplacian update which is similar to clustering.
    if L1.__len__() != L2.__len__():
        print('size(L1) must == size(L2)')

    Label1 = np.unique(L1)
    nClass1 = Label1.__len__()
    Label2 = np.unique(L2)
    nClass2 = Label2.__len__()

    nClass = max(nClass1, nClass2)
    G = np.zeros((nClass, nClass))
    for i in range(nClass1):
        for j in range(nClass2):
            G[i][j] = np.nonzero((L1 == Label1[i]) * (L2 == Label2[j]))[0].__len__()

    c = linear_sum_assignment(-G.T)[1]
    newL2 = np.zeros(L2.__len__())
    for i in range(nClass2):
        for j in np.nonzero(L2 == Label2[i])[0]:
            if len(Label1) > c[i]:
                newL2[j] = Label1[c[i]]

    return accuracy_score(L1, newL2),newL2


def lshot_prediction(args, knn, lmd, X, unary, support_label, test_label):
    # W = create_affinity(X, knn)
    W = None
    l, Y = bound_update(args, unary, W, lmd)
    out = np.take(support_label, l)
    out = out.astype(int)
    if args.lshot_remap:
        acc, _ = get_accuracy(test_label, out)
    else:
        acc = (test_label==out).mean()
    return acc, Y


def laplacian_shot_inference(args, gallery, query, support_label, test_label, train_mean=None, shot=5, lmd=.5):

    if args.train_mean is True:
        gallery = gallery - train_mean
        gallery = gallery / LA.norm(gallery, 2, 1)[:, None]
    if args.train_mean is True:
        query = query - train_mean
        query = query / LA.norm(query, 2, 1)[:, None]

    if args.proto_rect:
        eta = gallery.mean(0) - query.mean(0) # shift
        query = query + eta[np.newaxis,:]
        query_aug = np.concatenate((gallery, query),axis=0)
        gallery_ = gallery.reshape(args.meta_val_way, shot, gallery.shape[-1]).mean(1)
        gallery_ = torch.from_numpy(gallery_)
        query_aug = torch.from_numpy(query_aug)
        distance = get_metric('cosine')(gallery_, query_aug)
        predict = torch.argmin(distance, dim=1)
        cos_sim = F.cosine_similarity(query_aug[:, None, :], gallery_[None, :, :], dim=2)
        cos_sim = 10 * cos_sim
        W = F.softmax(cos_sim,dim=1)
        gallery_list = [(W[predict==i,i].unsqueeze(1)*query_aug[predict==i]).mean(0,keepdim=True) for i in sorted(predict.unique())]
        gallery = torch.cat(gallery_list,dim=0).numpy()
    else:
        gallery = gallery.reshape(args.meta_val_way, shot, gallery.shape[-1]).mean(1)

    support_label = support_label[::shot]
    subtract = gallery[:, None, :] - query
    distance = LA.norm(subtract, 2, axis=-1)
    test_label = np.array(test_label)
    # with LapLacianShot
    knn = args.knn
    lmd = lmd
    unary = distance.transpose() ** 2
    score, p = lshot_prediction(args, knn, lmd, query, unary, support_label, test_label)
    return score, p


def meta_evaluate(
            args, 
            tgt_test_out_dict, support_out_dict, 
            train_mean, 
            shot=5, lmd=.5
    ):
    (support_data, support_label,
     query_data, query_label) = meta_sample_data(
            tgt_test_out_dict, support_out_dict, 
            args.meta_val_way, args.meta_val_query, shot=shot
        )
    lap_score, Y = laplacian_shot_inference(
                args,
                support_data, query_data, 
                support_label, query_label, 
                train_mean=train_mean, 
                shot=shot, lmd=lmd
            )
    return lap_score, Y


def select_n_way_samples(args, processor, dataset, shot, way):
    # print("select_n_way_samples" , len(dataset))
    set_seed(args.seed)
    sampler = RandomSampler(copy.deepcopy(dataset)) 
    iterator = iter(DataLoader(copy.deepcopy(dataset), sampler=sampler, batch_size=4))
    inputs = {
        "input_ids": None,
        "attention_mask": None,
        "labels": None
    }

    num_class = len(processor.get_labels())
    class_dict = { i:0 for i in range(num_class) }
    cnt = 0
    while True:
        tot_sample = 0
        for k, v in class_dict.items():
            tot_sample += v
        if tot_sample == shot*way:
            for k, v in inputs.items():
                if k == "conf_penalty":
                    continue
                inputs[k].to(args.device)

            yield inputs
            inputs = {
                "input_ids": None,
                "attention_mask": None,
                "labels": None
            }
            class_dict = { i:0 for i in range(num_class) }
        try:
            curr_batch = iterator.__next__()
        except StopIteration:
            iterator = iter(DataLoader(copy.deepcopy(dataset), sampler=sampler, batch_size=4))
            curr_batch = iterator.__next__()

        curr_batch = tuple(t.to(args.device) for t in curr_batch)
        batch = curr_batch
        for _idx, class_id in enumerate(batch[3]):
            __int_class_id = class_id.item()
            if (class_dict[__int_class_id] + 1) <= shot:
                class_dict[__int_class_id] += 1
                if inputs["input_ids"] is None:
                    inputs["input_ids"] = torch.unsqueeze(batch[0][_idx], 0)
                    inputs["attention_mask"] = torch.unsqueeze(batch[1][_idx], 0)
                    inputs["labels"] = torch.unsqueeze(batch[3][_idx], 0)
                else:
                    inputs["input_ids"] = torch.cat((inputs["input_ids"], torch.unsqueeze(batch[0][_idx], 0)), 0)
                    inputs["attention_mask"] = torch.cat((inputs["attention_mask"], torch.unsqueeze(batch[1][_idx], 0)), 0)
                    inputs["labels"] = torch.cat((inputs["labels"], torch.unsqueeze(batch[3][_idx], 0)), 0)


def split_train_dev(args, support_inputs):
    if args.val_shot == 0:
        return support_inputs, support_inputs
    else:
        raise NotImplementedError


def get_score(task_name, logits, orig_preds):
    logits=logits.detach().cpu().numpy()
    eval_preds = np.argmax(logits, axis=1)
    orig_preds = orig_preds.detach().cpu().numpy()
    result = compute_metrics(task_name, eval_preds, orig_preds)
    return result


def get_metric(metric_type):
    METRICS = {
        'cosine': lambda gallery, query: 1. - F.cosine_similarity(query[:, None, :], gallery[None, :, :], dim=2),
        'euclidean': lambda gallery, query: ((query[:, None, :] - gallery[None, :, :]) ** 2).sum(2),
        'l1': lambda gallery, query: torch.norm((query[:, None, :] - gallery[None, :, :]), p=1, dim=2),
        'l2': lambda gallery, query: torch.norm((query[:, None, :] - gallery[None, :, :]), p=2, dim=2),
    }
    return METRICS[metric_type]


def batch_generator_from_sampler(support_inputs, batch_size=4):
    while True:
        ret_support_inputs = copy.deepcopy(support_inputs)
        rand_indices = torch.randperm(ret_support_inputs["input_ids"].size(0))
        for k, v in ret_support_inputs.items():
            if k == "sec_classifier":
                continue
            ret_support_inputs[k] = v[ rand_indices[0:batch_size] ]
        yield ret_support_inputs, rand_indices


def seqClassify_few_shot_extract_feature(args, model, src_ref_dataset, src_lang, force_load = False):
    model_hash = args.model_name_or_path.replace("/","_").replace("\\", "_")
    save_dir = os.path.join(args.output_dir, model_hash)
    os.makedirs(save_dir, exist_ok=True)
    model.eval()
    args.eval_batch_size = args.per_gpu_eval_batch_size
    with torch.no_grad():
        # get reference training mean
        if not os.path.isfile( os.path.join(save_dir,'{}.src_ref_dt_mean.plk'.format(src_lang)) ) or (force_load == True):
            out_mean, fc_out_mean = [], []
            sampler = SequentialSampler(src_ref_dataset) 
            src_ref_dataloader = DataLoader(src_ref_dataset, sampler=sampler, batch_size=args.eval_batch_size)
            for batch in src_ref_dataloader:
                batch = tuple(t.to(args.device) for t in batch)
                with torch.no_grad():
                    inputs = {
                                "input_ids": batch[0],
                                "attention_mask": batch[1],
                            }
                    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)
                    fc_outputs = outputs[1]
                    outputs = outputs[2][:, 0, :]
                    out_mean.append(outputs.cpu().data.numpy())
                    fc_out_mean.append(fc_outputs.cpu().data.numpy())

                    
            out_mean = np.concatenate(out_mean, axis=0).mean(0)
            fc_out_mean = np.concatenate(fc_out_mean, axis=0).mean(0)
            pickle.dump(
                    [out_mean, fc_out_mean], 
                    open(os.path.join(save_dir,'{}.src_ref_dt_mean.plk'.format(src_lang)), "wb")
            )
        else:
            out_mean, fc_out_mean = pickle.load(open(os.path.join(save_dir,'{}.src_ref_dt_mean.plk'.format(src_lang)), "rb"))
        src_ref_out_mean = out_mean
        src_ref_fc_out_mean = fc_out_mean
    return src_ref_out_mean, src_ref_fc_out_mean


def sample_case(
        tgt_test_out_dict, tgt_ref_out_dict, 
        tgt_test_fc_out_dict, tgt_ref_fc_out_dict, 
        meta_val_way, meta_val_query, shot
    ):
    # Sample meta task
    test_sample_class = sorted(random.sample(list(tgt_test_out_dict.keys()), meta_val_way))
#     ref_sample_class = random.sample(list(tgt_ref_out_dict.keys()), meta_val_way)
    ref_sample_class = test_sample_class
    train_input, test_input = [], []
    fc_train_input, fc_test_input = [], []
    test_label, train_label = [], []
#     print("test_sample_class", test_sample_class)
#     print("ref_sample_class", ref_sample_class)
#     print("meta_val_query", meta_val_query)
#     print("shot", shot)
    for each_class in ref_sample_class: 
        total_samples = shot
        if len(tgt_ref_out_dict[each_class]) < total_samples:
            total_samples = len(tgt_ref_out_dict[each_class])
#         samples = random.sample(tgt_ref_out_dict[each_class] , total_samples)
#         print("----------------:", samples)
        sample_list = random.sample(list(zip(tgt_ref_out_dict[each_class],tgt_ref_fc_out_dict[each_class])) , total_samples)
        samples, fc_samples = [], []
        for a, b in sample_list:
            samples += [a]
            fc_samples += [b]
        train_label += [each_class] * len(samples)
        train_input += samples
        fc_train_input += fc_samples
    for each_class in test_sample_class: 
        total_samples = meta_val_query
        if len(tgt_test_out_dict[each_class]) < total_samples:
            total_samples = len(tgt_test_out_dict[each_class])
        sample_list = random.sample(list(zip(tgt_test_out_dict[each_class], tgt_test_fc_out_dict[each_class])), total_samples)
        samples, fc_samples = [], []
        for a, b in sample_list:
            samples += [a]
            fc_samples += [b]
        test_label += [each_class] * len(samples)
        test_input += samples
        fc_test_input += fc_samples
    train_input = np.array(train_input).astype(np.float32)
    test_input = np.array(test_input).astype(np.float32)
    return train_input, test_input, fc_train_input, fc_test_input, train_label, test_label


def dev_meta_evaluate(args, 
        tgt_test_out_dict, tgt_ref_out_dict, 
        tgt_test_fc_out_dict, tgt_ref_fc_out_dict, 
        src_ref_out_mean, 
        shot=5, lmd=.5
    ):
    un_list = []
    l2n_list = []
    cl2n_list = []
    baseline_cl2n_list = []
    # for _ in warp_tqdm(True, "lshot", range(args.meta_test_iter)):
    for _ in range(args.meta_test_iter):
        (train_data, test_data, 
         fc_train_input, fc_test_input,
         train_label, test_label) = sample_case(
                tgt_test_out_dict, tgt_ref_out_dict, 
                tgt_test_fc_out_dict, tgt_ref_fc_out_dict, 
                args.meta_val_way, args.meta_val_query, shot
         )
        correct, total = 0, 0
        for fc_out, label in zip(fc_test_input, test_label) :
            pred_label = np.argmax(fc_out, axis=-1)
            if pred_label == label:
                correct += 1
            total += 1
        baseline_acc =  correct/float(total)
#         baseline_acc = 0
        baseline_cl2n_list.append(baseline_acc*100)
        acc = laplacian_shot_inference(
                args,
                train_data, test_data, 
                train_label, test_label, 
                train_mean=src_ref_out_mean, 
                shot=shot, lmd=lmd
            )
        cl2n_list.append(acc*100)
    cl2n_mean, cl2n_conf = compute_confidence_interval(cl2n_list)
    baseline_cl2n_mean, baseline_cl2n_conf = compute_confidence_interval(baseline_cl2n_list)
    return cl2n_mean, cl2n_conf, baseline_cl2n_mean, baseline_cl2n_conf


def get_zero_shot_score(args, model, eval_inputs):
    with torch.no_grad():
        _, (_, zero_shot_eval_logits, zero_shot_eval_seq_outputs) = model(**eval_inputs)
        zero_shot_result = get_score(args.task_name, zero_shot_eval_logits, eval_inputs["labels"])
    score = zero_shot_result[args.model_selection_metric]
    return score, zero_shot_eval_logits, zero_shot_eval_seq_outputs

def knn_data_prepare(args, model, support_inputs, eval_inputs, zero_shot_eval_logits=None, zero_shot_eval_seq_outputs=None):
    
    model.eval()
    with torch.no_grad():
        _, (_, zero_shot_support_logits, zero_shot_support_seq_output) = model(**support_inputs)
    
    tgt_test_out_dict = collections.defaultdict(list)
    tgt_test_fc_out_dict = collections.defaultdict(list)
    
    if zero_shot_eval_logits is None or zero_shot_eval_seq_outputs is None:
        (_, 
         zero_shot_eval_logits, 
         zero_shot_eval_seq_outputs) = get_zero_shot_score(
                args, model, eval_inputs
            )
        
    fc_outputs = zero_shot_eval_logits
    outputs = zero_shot_eval_seq_outputs[:, 0, :]
    outputs = outputs.cpu().data.numpy()
    fc_outputs = fc_outputs.cpu().data.numpy()
    for out, fc_out, label in zip(outputs, fc_outputs, eval_inputs["labels"]):
        tgt_test_out_dict[label.item()].append(out)
        tgt_test_fc_out_dict[label.item()].append(fc_out)
    
    tgt_ref_out_dict = collections.defaultdict(list)
    tgt_ref_fc_out_dict = collections.defaultdict(list)
    fc_outputs = zero_shot_support_logits
    outputs = zero_shot_support_seq_output[:, 0, :]
    outputs = outputs.cpu().data.numpy()
    fc_outputs = fc_outputs.cpu().data.numpy()
    for out, fc_out, label in zip(outputs, fc_outputs, support_inputs["labels"]):
        tgt_ref_out_dict[label.item()].append(out)
        tgt_ref_fc_out_dict[label.item()].append(fc_out)
    return tgt_test_out_dict, tgt_test_fc_out_dict, tgt_ref_out_dict, tgt_ref_fc_out_dict

def get_knn_score(args, tgt_test_out_dict, tgt_ref_out_dict, train_mean):
    knn_acc, Y = meta_evaluate(
                args, 
                tgt_test_out_dict, tgt_ref_out_dict, 
                train_mean, 
                shot=args.shot, lmd=0.0
            )
    return knn_acc, Y
        
    
def log_stat(key, values, episode, logger=None, tb_writer=None):
    acc_lst = []
    if key.startswith("zero_shot"):
        for val in values:
            acc_lst.append(val[0])
    elif key.startswith("knn"):
        for val in values:
            (knn_acc, P_knn) , exec_time = val 
            acc_lst.append(knn_acc)
    elif key.startswith("finetuning"):
        for val in values:
            ((score, epoch_scores), exec_time) = val 
            acc_lst.append(score)
    else:
        raise NotImplementedError()
    score = round(np.average(acc_lst)*100,2)
    if logger is not None:
        logger.info("\t\t\t{} : {}".format(key, score))
    else:
        print("\t\t\t{} : {}".format(key, score))
    if tb_writer is not None:
        tb_writer.add_scalar(key, score, episode)

def get_source_idx(args):
    src_idx = -1
    for idx, file_info in enumerate(args.dev):
        if file_info.split(";")[-1] == args.src_lang:
            src_idx = idx
            break
    return src_idx

def SequenceClassificationEvaluateLshotBenchmark(
        args,
        base_model,
        tokenizer,
        processor,
        support_file_info,
        test_file_info,
        conf_penalty=0,
        logger=None,
        verbose=True
    ):

    if logger is None:
        logger = logging.getLogger(__name__)
    if verbose:
        logger.info("Evaluating ...")
        logger.info("<"+"-"*10+">")
        logger.info("  Support file : {}".format(support_file_info))
        logger.info("  Test file : {}".format(test_file_info))

    # load datasets
    eval_dataset = processor.load_and_cache_examples(
            args, tokenizer, mode='test', data_file_info = test_file_info, logger= logger
        )
    support_dataset = processor.load_and_cache_examples(
                    args, tokenizer, mode='support', data_file_info = support_file_info, logger= logger
        )

    src_idx = get_source_idx(args)
    source_dataset = processor.load_and_cache_examples(
                    args, tokenizer, mode='support', data_file_info = args.dev[src_idx], logger = logger
        )


    train_mean, support_fc_out_mean = None, None
    if "knn" in args.benchmarks:
        prev_base_model_device = base_model.device
        base_model.to(args.device)
        train_mean, support_fc_out_mean = seqClassify_few_shot_extract_feature(
                args, base_model, source_dataset, args.src_lang,
                force_load = True
            )
        base_model.to(prev_base_model_device)
    
    # model
    prev_base_model_device = base_model.device
    base_model.to(args.device)
    base_model.to(prev_base_model_device)
    
    eval_iter = select_n_way_samples(
            args, processor, eval_dataset, 
            shot=args.meta_val_query, way=args.meta_val_way
        )
    support_iter = select_n_way_samples(
            args, processor, support_dataset, 
            shot=args.shot, way=args.meta_val_way
        )
    
    logger.info("Episoding Testing started ...")
    result_dict = {
        "zero_shot":[], 
    }
    
    if args.local_rank in [-1, 0]:
        # tensorboard setup
        tb_path = os.path.join(args.output_dir, "tf_board")
        if not os.path.exists(tb_path):
            os.makedirs(tb_path)
        tb_writer = SummaryWriter(tb_path)

    iteration_interval_start_time = time.time()
    iteration_strat_time = iteration_interval_start_time
    iteration_strat_episode = 0
    for episode in range(args.meta_test_iter):

        support_inputs = support_iter.__next__()
        eval_inputs = eval_iter.__next__()  
        if args.overwrite_num_of_label is not None:
            support_inputs["sec_classifier"] = True 
        if args.overwrite_num_of_label is not None:
            eval_inputs["sec_classifier"] = True 
        
        model = copy.deepcopy(base_model)
        model.to(args.device)

        # load eval data
        model.eval()
        zero_shot_eval_logits = None
        zero_shot_eval_seq_outputs = None
        # Zero Shot Accuracy
        if "zero_shot" in args.benchmarks: 
            start_time = time.time()
            (zs_score, 
             zero_shot_eval_logits, 
             zero_shot_eval_seq_outputs) = get_zero_shot_score(
                    args, model, eval_inputs
                )
            zero_shot_time = time.time() - start_time
            result_dict["zero_shot"].append((zs_score, zero_shot_time))
        # KNN Accuracy
        if "knn" in args.benchmarks: 
            start_time = time.time()
            (tgt_test_out_dict, tgt_test_fc_out_dict, 
            tgt_ref_out_dict, tgt_ref_fc_out_dict) = knn_data_prepare(
                        args, model, 
                        support_inputs, eval_inputs, 
                        zero_shot_eval_logits, zero_shot_eval_seq_outputs
                    )
            data_prep_time = time.time() - start_time
            # data prep time can be precalculated
            for norm in [0, 1]:
                if norm == 0:
                    args.train_mean = False
                elif norm == 1:
                    args.train_mean = True
                for proto_rect in [0, 1]:
                    args.proto_rect=proto_rect
                    start_time = time.time()
                    knn_acc, _ = get_knn_score(
                            args, tgt_test_out_dict, tgt_ref_out_dict, train_mean
                        )
                    knn_infer_prep_time = time.time() - start_time
                    task_name = "knn"
                    task_name = "knn_norm" if args.train_mean == True else "knn"
                    task_name = "{}{}".format(task_name, "" if proto_rect == 0 else "_proto_rect")
                    if task_name not in result_dict:
                        result_dict[task_name] = []
                    result_dict[task_name].append(
                        ( (knn_acc, data_prep_time) , knn_infer_prep_time )
                    )
        model = None
        if "finetuning" in args.benchmarks: 
            support_inputs, dev_inputs = split_train_dev(args, support_inputs)
            for fs_finetune_lr in args.fs_finetune_lr:
                for fs_finetune_batch in args.fs_finetune_batch:
                    for fs_grad_acc_step in args.fs_grad_acc_step:
                        for fs_finetune_lr_scheduler in args.fs_finetune_lr_scheduler:
                            for fs_finetune_type in args.fs_finetune_type:
                                for fs_finetune_cf in args.fs_finetune_cf:
                                    for fs_finetune_transductive in args.fs_finetune_transductive:
                                        if fs_finetune_transductive == 2 and fs_finetune_type == "full":
                                            continue
                                        exp_name = "finetuning_lr_{}-lrsc_{}-bcz_{}-fnt_{}-cf_{}-trans_{}".format(
                                            fs_finetune_lr, fs_finetune_lr_scheduler, fs_finetune_batch*fs_grad_acc_step,
                                            fs_finetune_type, fs_finetune_cf, fs_finetune_transductive
                                        )
                                        if exp_name not in result_dict:
                                            result_dict[exp_name] = []
                                        trainer = EpisodicFinetuneTrainer(
                                            model=copy.deepcopy(base_model),
                                            args=args,
                                            epoch=args.fewshot_benchmark_epoch,
                                            lr_rate=fs_finetune_lr,
                                            batch_size=fs_finetune_batch,
                                            grad_acc_step=fs_grad_acc_step,
                                            scheduler_type=fs_finetune_lr_scheduler,
                                            optimization=fs_finetune_type,
                                            conf_penalty=fs_finetune_cf,
                                            transduction=fs_finetune_transductive
                                        )
                                        if fs_finetune_transductive > 0:
                                            trainer.reset_optimizer_scheduler(optimization_type="head")
                                            ((score, epoch_score), exec_time1)=  trainer.train(
                                                support_inputs, dev_inputs, eval_inputs
                                            )
                                            ((score, epoch_score), exec_time2) =  trainer.train(
                                                support_inputs, dev_inputs, eval_inputs, is_transduction=True, P_knn=P_knn
                                            )
                                            exec_time = exec_time1 + exec_time2
                                        else:
                                            trainer.reset_optimizer_scheduler(optimization_type=fs_finetune_type)
                                            ((score, epoch_score), exec_time)=  trainer.train(
                                                support_inputs, dev_inputs, eval_inputs
                                            )
                                        result_dict[exp_name].append(
                                            ((score, epoch_score), exec_time)
                                        )

        # test num of times eval
        if episode <= 10 or (episode > 10 and episode%10 == 0):
            current_time = time.time()
            iteration_interval = current_time - iteration_interval_start_time
            total_episode_done = episode - iteration_strat_episode + 1
            curr_speed = round(iteration_interval/total_episode_done, 1)
            total_elapsed_time = current_time - iteration_strat_time
            total_completed_episode = episode+1
            remaining_episode = args.meta_test_iter-total_completed_episode
            est_time = round(((total_elapsed_time/total_completed_episode)*remaining_episode))
            est_time = time.strftime('%H:%M:%S', time.gmtime(est_time))
            
            logger.info("Logging Resutls at episode : {}, Speed : {} (sec/episode), Estimated Time Needed: {}".format(episode, curr_speed, est_time))
            for k, v in result_dict.items():
                log_stat(k, v, episode, logger, tb_writer)
            
            iteration_interval_start_time = time.time()
            iteration_strat_episode = episode + 1

    return result_dict


class EpisodicFinetuneTrainer(object):

    def __init__(self, model, args, epoch, lr_rate, batch_size, grad_acc_step, scheduler_type, optimization, conf_penalty, transduction):
        self.model=model
        self.args=args
        self.epoch=epoch
        self.lr_rate=lr_rate
        self.batch_size=batch_size
        self.grad_acc_step=grad_acc_step
        self.scheduler_type=scheduler_type
        self.optimization=optimization
        self.conf_penalty=conf_penalty
        self.transduction=transduction
        self.reset_optimizer_scheduler(optimization_type=self.optimization)


    def reset_optimizer_scheduler(self, optimization_type=None):
        self.optimizer = get_optimizer(
                    self.model, 
                    self.args.weight_decay, 
                    self.args.learning_rate, 
                    self.args.adam_epsilon,
                    optimization=optimization_type,
                    is_sec_classifier=True if self.args.overwrite_num_of_label is not None else False
                )
        
        self.total_training_samples = (self.args.meta_val_way * self.args.shot - self.args.meta_val_way*self.args.val_shot)*self.epoch
        self.total_training_steps = int(round((self.total_training_samples/(self.batch_size*self.grad_acc_step))+1))
        self.eval_step = int(self.total_training_samples/self.epoch)
        if self.scheduler_type == "linear_decay":
            self.scheduler = get_linear_schedule_with_warmup(
                    self.optimizer, 
                    num_warmup_steps=0, 
                    num_training_steps=self.total_training_steps
                )
        else:
            self.scheduler = get_constant_schedule(self.optimizer)


    def train(self, support_inputs, val_inputs, eval_inputs, is_transduction=False, P_knn=None):
        
        start_time = time.time()

        if is_transduction == True:
            self.total_training_samples += (self.args.meta_val_query * self.args.meta_val_way)*self.epoch
            self.total_training_steps = int(round((self.total_training_samples/(self.batch_size*self.grad_acc_step))+1))
            self.eval_step = int(self.total_training_samples/self.epoch)
        
        self.model.to(self.args.device)
        self.model.train()
        tune_time = 0
        if is_transduction:
            tgt_data_iterator = batch_generator_from_sampler(eval_inputs, batch_size=int(self.batch_size*0.75))
            src_data_iterator = batch_generator_from_sampler(support_inputs, batch_size=int(self.batch_size*0.25))
        else:
            src_data_iterator = batch_generator_from_sampler(support_inputs, batch_size=int(self.batch_size))
        
        mx_sup_val_res, max_sup_res = -1, -1
        samples_cnt, eval_cnt = 0, 0
        global_step = 0
        for step in range(self.total_training_steps * self.grad_acc_step):
            self.model.train()
            if is_transduction:
                dt, index = tgt_data_iterator.__next__()
                dt["conf_penalty"] = True 
                dt["marginal_entropy"] = True
                dt["labels"] = None 
                dt["conf_coef"] = .1 
                if self.args.overwrite_num_of_label is not None:
                    dt["sec_classifier"] = True 
                mi_loss, (_, _, _) = self.model(**dt)
                mi_loss = mi_loss
                samples_cnt += int(dt['input_ids'].size()[0])

            dt, index = src_data_iterator.__next__()
            if self.args.overwrite_num_of_label is not None:
                dt["sec_classifier"] = True 
                
            loss, (_, _, _) = self.model(**dt)
            
            samples_cnt += int(dt['input_ids'].size()[0])

            total_loss = loss 
            if is_transduction:
                total_loss -= mi_loss
            
            total_loss=total_loss/self.grad_acc_step
            total_loss.backward()
            if (step+1) % self.grad_acc_step == 0:
                self.optimizer.step()
                self.scheduler.step()
                self.model.zero_grad()
                global_step += 1
            
            if samples_cnt >= self.eval_step:
                tune_time_start = time.time()
                eval_cnt += 1
                samples_cnt = 0
                self.model.eval()
                with torch.no_grad():
                    if self.args.overwrite_num_of_label is not None:
                        val_inputs["sec_classifier"] = True 
                    _, (_, val_inputs_logits, _) = self.model(**val_inputs)
                    val_inputs_result = get_score(self.args.task_name, val_inputs_logits, val_inputs['labels'])
                if val_inputs_result[self.args.model_selection_metric] > mx_sup_val_res :
                    mx_sup_val_res = val_inputs_result[self.args.model_selection_metric]
                    with torch.no_grad():
                        if self.args.overwrite_num_of_label is not None:
                            eval_inputs["sec_classifier"] = True 
                        _, (_, eval_logits, _) = self.model(**eval_inputs)
                        supervised_result = get_score(self.args.task_name, eval_logits, eval_inputs['labels'])
                        max_sup_res = supervised_result[self.args.model_selection_metric]
                tune_time += (time.time() - tune_time_start)

        tune_time_start = time.time()
        eval_cnt += 1
        self.model.eval()
        with torch.no_grad():
            if self.args.overwrite_num_of_label is not None:
                val_inputs["sec_classifier"] = True 
            _, (_, val_inputs_logits, _) = self.model(**val_inputs)
            val_inputs_result = get_score(self.args.task_name, val_inputs_logits, val_inputs['labels'])
        tune_time += (time.time() - tune_time_start)

        with torch.no_grad():
            if self.args.overwrite_num_of_label is not None:
                eval_inputs["sec_classifier"] = True 
            _, (_, eval_logits, _) = self.model(**eval_inputs)
            supervised_result = get_score(self.args.task_name, eval_logits, eval_inputs['labels'])
        without_finetuning_res = supervised_result[self.args.model_selection_metric]
        if val_inputs_result[self.args.model_selection_metric] > mx_sup_val_res :
            mx_sup_val_res = val_inputs_result[self.args.model_selection_metric]
            max_sup_res = supervised_result[self.args.model_selection_metric]
        
        exec_time = time.time() - start_time
        return ((max_sup_res,[without_finetuning_res, tune_time]), exec_time)