test.py

import os
import json
import torch
import numpy
import logging
import random
from tqdm import tqdm
from torch import Tensor
from torch.nn import functional as F
from collections import defaultdict
from utils.dataloader import make_loader


def top_k_top_p_filtering(
        logits: Tensor,
        top_k: int = 0,
        top_p: float = 1.0,
        filter_value: float = -float("Inf"),
        min_tokens_to_keep: int = 1,
) -> Tensor:
    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
    Args:
        logits: logits distribution shape (batch size, vocabulary size)
        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
        Make sure we keep at least min_tokens_to_keep per batch example in the output
    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
    """
    if top_k > 0:
        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
        # Remove all tokens with a probability less than the last token of the top-k
        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
        logits[indices_to_remove] = filter_value

    if top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
        sorted_indices_to_remove = cumulative_probs > top_p
        if min_tokens_to_keep > 1:
            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
        # Shift the indices to the right to keep also the first token above the threshold
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        # scatter sorted tensors to original indexing
        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
        logits[indices_to_remove] = filter_value
    return logits


def get_example_inputs(model, tokenizer, prompt_text, device):
    num_attention_heads = model.config.n_head
    hidden_size = model.config.n_embd
    num_layer = model.config.n_layer
    tokenizer.padding_side = "left"
    encodings_dict = tokenizer.batch_encode_plus(prompt_text, padding=True)

    input_ids = torch.tensor(encodings_dict['input_ids'], dtype=torch.int64)
    attention_mask = torch.tensor(encodings_dict['attention_mask'], dtype=torch.float32)
    position_ids = (attention_mask.long().cumsum(-1) - 1)
    position_ids.masked_fill_(position_ids < 0, 0)

    # Empty Past State for generating first word
    empty_past = []
    batch_size = input_ids.size(0)
    sequence_length = input_ids.size(1)
    past_shape = [2, batch_size, num_attention_heads, 0, hidden_size // num_attention_heads]
    for i in range(num_layer):
        empty_past.append(torch.empty(past_shape).type(torch.float32).to(device))

    return input_ids.to(device), attention_mask.to(device), position_ids.to(device), empty_past


def test_generation_GPT2BATCH(model, tokenizer, input_text, device, do_sample=False, temperature=1.0, top_k=0, top_p=0,
                              max_length=30, task_id=-1):
    eos_token_id = tokenizer.eos_token_id

    input_ids, attention_mask, position_ids, past = get_example_inputs(model, tokenizer, input_text, device)
    batch_size = input_ids.size(0)

    has_eos = torch.zeros(batch_size, dtype=torch.bool).to(device)

    all_token_ids = input_ids.clone()

    for step in range(max_length):

        if task_id == -1:
            outputs = model(input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past,
                            return_dict=False)
        else:
            outputs = model(input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past,
                            task_id=task_id, return_dict=False)

        next_token_logits = outputs[0][:, -1, :]

        if do_sample:
            # Temperature (higher temperature => more likely to sample low probability tokens)
            if temperature != 1.0:
                next_token_logits = next_token_logits / temperature
            # Top-p/top-k filtering
            next_token_logscores = top_k_top_p_filtering(next_token_logits, top_k=top_k, top_p=top_p)
            # Sample
            probs = F.softmax(next_token_logscores, dim=-1)
            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
        else:
            # Greedy decoding
            next_tokens = torch.argmax(next_token_logits, dim=-1)

        has_eos = has_eos | (next_tokens == eos_token_id)
        tokens_to_add = next_tokens.masked_fill(has_eos, eos_token_id)
        all_token_ids = torch.cat([all_token_ids, tokens_to_add.unsqueeze(-1)], dim=-1)

        # Update input_ids, attention_mask, position_ids and past
        input_ids = tokens_to_add.clone().detach().reshape([batch_size, 1]).to(device)
        position_ids = (position_ids[:, -1] + 1).reshape(batch_size, 1)
        attention_mask = torch.cat([attention_mask, torch.ones([batch_size, 1]).type_as(attention_mask)], 1).to(device)

        past = list(outputs[1])  # past in torch output is tuple
        if torch.all(has_eos):
            break

    responses = []
    responses_plain = []
    for i, output in enumerate(all_token_ids):
        responses_plain.append(tokenizer.decode(output, skip_special_tokens=True))
        res = tokenizer.decode(output, skip_special_tokens=True)
        responses.append(res[res.find("[SOS]"):].replace("[SOS]", "").strip())
    return responses, responses_plain


def generate_sample_prev_task(args, model, tokenizer, dataset_dic, task_id_so_far, number_of_sample, time,
                              task_id_adpt=-1):
    # device = torch.device(f"cuda:{args.GPU[0]}")
    device = torch.device(f"cuda:0")
    model.to(device)
    model.eval()
    ## notice that this sample is just to have the data struct
    temp_aug_mem = random.sample(dataset_dic["['sgd_restaurants']"],
                                 min(len(dataset_dic["['sgd_restaurants']"]), number_of_sample))
    temp_aug_sam = random.sample(dataset_dic["['sgd_restaurants']"],
                                 min(len(dataset_dic["['sgd_restaurants']"]), number_of_sample))
    with torch.no_grad():
        if "gpt2" in args.model_checkpoint:  ## this works only with GPT2
            sample_list = []
            for i in range(int(number_of_sample / (args.valid_batch_size)) + 1):
                if (i % 2 == 0 or args.task_type != "E2E"):  # sample on batch with and one without API call
                    input_batch = [f"[{str(eval(task_id_so_far)[0])}]" for _ in range(args.valid_batch_size)]
                else:
                    input_batch = [f"[{str(eval(task_id_so_far)[0])}-API]" for _ in range(args.valid_batch_size)]
                _, samples = test_generation_GPT2BATCH(model=model,
                                                       tokenizer=tokenizer,
                                                       input_text=input_batch,
                                                       device=device,
                                                       max_length=300,
                                                       do_sample=True,
                                                       top_p=0.9,
                                                       temperature=1.1,
                                                       task_id=task_id_adpt)
                sample_list += samples
    sample_list = random.sample(sample_list, min(len(sample_list), number_of_sample))
    # this sample is to train the previous task generator
    for i in range(len(temp_aug_mem)):
        temp_aug_mem[i]["history_reply"] = f"{sample_list[i].strip()} {tokenizer.eos_token}"
    # this sample is to train the previous task itself
    # hence we remove the special token in input
    for i in range(len(temp_aug_sam)):
        samp = sample_list[i].strip()
        samp = samp.replace(f"[{str(eval(task_id_so_far)[0])}]", "")
        samp = samp.replace(f"[{str(eval(task_id_so_far)[0])}-API]", "")
        temp_aug_sam[i]["history_reply"] = f"{samp} {tokenizer.eos_token}"

    temp_aug = temp_aug_mem + temp_aug_sam
    ## save the generated data for logging
    if not os.path.exists(f'{args.saving_dir}/{time}'):
        os.makedirs(f'{args.saving_dir}/{time}')
    with open(f'{args.saving_dir}/{time}' + f'/{task_id_so_far}_generated.json', 'w') as fp:
        json.dump(temp_aug, fp, indent=4)
    return temp_aug


def test_model_seq2seq(args, model, tokenizer, test_loader, time="0_['']"):
    device = torch.device(f"cuda:0")
    model.to(device)
    model.eval()
    results = []

    for idx_b, batch in tqdm(enumerate(test_loader), total=len(test_loader)):
        with torch.no_grad():
            if "gpt2" in args.model_checkpoint:
                value_batch, _ = test_generation_GPT2BATCH(model=model,
                                                           tokenizer=tokenizer,
                                                           input_text=[b + "[SOS]" for b in batch['history']],
                                                           device=device,
                                                           max_length=100)
            else:
                responses = model.generate(input_ids=batch["encoder_input"].to(device),
                                           attention_mask=batch["attention_mask"].to(device),
                                           eos_token_id=tokenizer.eos_token_id,
                                           max_length=100)
                value_batch = tokenizer.batch_decode(responses, skip_special_tokens=True)
        for idx, resp in enumerate(value_batch):
            results.append({"id": batch["dial_id"][idx], "turn_id": batch["turn_id"][idx],
                            "dataset": batch["dataset"][idx], "task_id": batch["task_id"][idx],
                            "spk": batch["spk"][idx], "gold": batch["reply"][idx],
                            "genr": resp, "hist": batch["history"][idx]})
        # if(idx_b==1): break
    if not os.path.exists(f'{args.saving_dir}/{time}'):
        os.makedirs(f'{args.saving_dir}/{time}')
    with open(f'{args.saving_dir}/{time}' + '/generated_responses.json', 'w') as fp:
        json.dump(results, fp, indent=4)
    tokenizer.padding_side = "right"


def argmin(a):
    return min(range(len(a)), key=lambda x: a[x])


def test_model_seq2seq_ADAPTER(args, model, tokenizer, test_loader, test_dataset, time="0_['']"):
    # device = torch.device(f"cuda:{args.GPU[0]}")
    device = torch.device(f"cuda:0")
    model.model.to(device)
    model.model.eval()
    results = []

    print(model.task_list_seen, len(model.task_list_seen))
    # range_adpt = len(model.task_list_seen)

    # perplexity_dict = {f'{sample["dial_id"]}_{sample["turn_id"]}_{sample["task_id"]}': [] for sample in test_dataset}
    # for t in range(range_adpt):
    #     print(f"Task {t}")
    #     for idx_b, batch in tqdm(enumerate(test_loader), total=len(test_loader)):
    #         ppl_batch = model.compute_PPL(batch, task_id=t, device=device)  ## one value per batch
    #         for (d_id, t_id, ta_id, ppl) in zip(batch["dial_id"], batch["turn_id"], batch["task_id"], ppl_batch):
    #             perplexity_dict[f'{d_id}_{t_id}_{ta_id}'].append(ppl)

    # select the task id with the lowest perplexity (loss)
    # perplexity_dict_ = {}
    # for k, v in perplexity_dict.items():
    #     if len(v) == range_adpt:
    #         perplexity_dict_[k] = v
    #     else:
    #         print(k, v)
    #
    # perplexity_dict = {k: argmin(v) for k, v in perplexity_dict_.items()}
    #
    # ## group by sample by predicted task id
    # test_dataset_by_predicted_id = defaultdict(list)
    # for sample in test_dataset:
    #     if (f'{sample["dial_id"]}_{sample["turn_id"]}_{sample["task_id"]}' in perplexity_dict):
    #         test_dataset_by_predicted_id[
    #             perplexity_dict[f'{sample["dial_id"]}_{sample["turn_id"]}_{sample["task_id"]}']].append(sample)
    #
    # for k, v in test_dataset_by_predicted_id.items():
    #     print(f"Task {k}: {len(v)}")
    #
    # ## create a dataloader for batch each of this
    # test_dataset_by_predicted_id = {k: make_loader(args, v, model.tokenizer) for k, v in
    #                                 test_dataset_by_predicted_id.items()}

    for task_id, task_loader in test_loader.items():
        pred_task_id = model.task_list_seen.index(task_id)
        # print(f"Task Id: {task_id}")
        for idx_b, batch in tqdm(enumerate(task_loader), total=len(task_loader), desc=task_id):
            with torch.no_grad():
                value_batch, _ = test_generation_GPT2BATCH(model=model.model,
                                                           tokenizer=model.tokenizer,
                                                           input_text=[b + "[SOS]" for b in batch['history']],
                                                           device=device,
                                                           max_length=100,
                                                           task_id=pred_task_id)
            for idx, resp in enumerate(value_batch):
                results.append({"id": batch["dial_id"][idx], "turn_id": batch["turn_id"][idx],
                                "dataset": batch["dataset"][idx], "task_id": batch["task_id"][idx],
                                "spk": batch["spk"][idx], "gold": batch["reply"][idx],
                                "genr": resp, "hist": batch["history"][idx],
                                "pred_task_id": model.task_list_seen[pred_task_id]})

            # if(idx_b==1): break
    if not os.path.exists(f'{args.saving_dir}/{time}'):
        os.makedirs(f'{args.saving_dir}/{time}')
    with open(f'{args.saving_dir}/{time}' + '/generated_responses.json', 'w') as fp:
        json.dump(results, fp, indent=4)
    tokenizer.padding_side = "right"


def test_model_seq2seq_ADAPTER_FWT(args, model, tokenizer, test_loader, test_dataset, time="0_['']"):
    # device = torch.device(f"cuda:{args.GPU[0]}")
    device = torch.device(f"cuda:0")
    model.model.to(device)
    model.model.eval()
    results = []

    print(model.task_list_seen, len(model.task_list_seen))

    for task_id, task_loader in test_loader.items():
        pred_task_id = model.task_list_seen.index(task_id)

        if pred_task_id == 0:
            continue
        # print(f"Task Id: {task_id}")
        for idx_b, batch in tqdm(enumerate(task_loader), total=len(task_loader), desc=task_id):
            with torch.no_grad():
                value_batch, _ = test_generation_GPT2BATCH(model=model.model,
                                                           tokenizer=model.tokenizer,
                                                           input_text=[b + "[SOS]" for b in batch['history']],
                                                           device=device,
                                                           max_length=100,
                                                           task_id=pred_task_id)
            for idx, resp in enumerate(value_batch):
                results.append({"id": batch["dial_id"][idx], "turn_id": batch["turn_id"][idx],
                                "dataset": batch["dataset"][idx], "task_id": batch["task_id"][idx],
                                "spk": batch["spk"][idx], "gold": batch["reply"][idx],
                                "genr": resp, "hist": batch["history"][idx],
                                "pred_task_id": model.task_list_seen[pred_task_id]})

            # if(idx_b==1): break
    if not os.path.exists(f'{args.saving_dir}/{time}'):
        os.makedirs(f'{args.saving_dir}/{time}')
    with open(f'{args.saving_dir}/{time}' + '/generated_responses.json', 'w') as fp:
        json.dump(results, fp, indent=4)
    tokenizer.padding_side = "right"



def generate_t5_adapter(model, tokenizer, batch, device, max_length=100, task_id=-1):
    model.to(device)
    outputs = model.generate(
        input_ids=batch['encoder_input'].to(device),
        attention_mask=batch['attention_mask'].to(device),
        use_cache=False,
        return_dict_in_generate=True,
        max_length=max_length,
        task_id=task_id,
    )

    responses = []
    for i, gen_ids in enumerate(outputs.sequences):
        res = tokenizer.decode(gen_ids, skip_special_tokens=True)
        responses.append(res)
    return responses


def test_t5_adapter(model, test_loader, args, time):
    device = torch.device(f"cuda:0")
    model.model.to()
    model.model.eval()
    results = []

    # print(model.task_list_seen, len(model.task_list_seen))

    for task_id, task_loader in test_loader.items():
        # print(f"Task Id: {task_id}")
        task_num_id = model.task_list_seen.index(task_id)
        for idx_b, batch in tqdm(enumerate(task_loader), total=len(task_loader)):
            with torch.no_grad():
                value_batch = generate_t5_adapter(model=model.model,
                                                  tokenizer=model.tokenizer,
                                                  batch=batch,
                                                  device=device,
                                                  max_length=100,
                                                  task_id=task_num_id)
            for idx, resp in enumerate(value_batch):
                results.append({"id": batch["dial_id"][idx], "turn_id": batch["turn_id"][idx],
                                "dataset": batch["dataset"][idx], "task_id": batch["task_id"][idx],
                                "spk": batch["spk"][idx], "gold": batch["reply"][idx],
                                "genr": resp, "hist": batch["history"][idx], "pred_task_id": task_id})

    if not os.path.exists(f'{args.saving_dir}/{time}'):
        os.makedirs(f'{args.saving_dir}/{time}')
    with open(f'{args.saving_dir}/{time}' + '/generated_responses.json', 'w') as fp:
        json.dump(results, fp, indent=4)


# def test_model(args,model,tokenizer,test_loader,time=0):
# args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# model.to(args.device)
# model.eval()
# results = []
# for task_id, task_loader in tqdm(test_loader.items(),total=len(test_loader)):
#     # print(f"Task Id: {task_id}")
#     for idx_b, batch in enumerate(task_loader):
#         input_ids, _, token_type_ids  = tuple(torch.tensor([batch[input_name]]).to(args.device) for input_name in MODEL_INPUTS)
#         with torch.no_grad():
#             response = generate(args,model,tokenizer,input_ids,token_type_ids)
#         results.append({"id":batch["dial_id"],"turn_id":batch["turn_id"],
#                         "dataset":batch["dataset"],"task_id":task_id,
#                         "spk":batch["spk"],"gold":batch["row_reply"],
#                         "genr":response,"hist":batch["plain_history"]})

# if not os.path.exists(f'{args.saving_dir}/{time}'):
#     os.makedirs(f'{args.saving_dir}/{time}')
# with open(f'{args.saving_dir}/{time}'+'/generated_responses.json', 'w') as fp:
#     json.dump(results, fp, indent=4)


def test():
    pass
    # args = get_args()
    # model = Seq2SeqToD(args)
    # model.model.load_state_dict(torch.load(f'runs_INTENT/BEST/ADAPTER_EPC_10_LR_0.00625_BOTL_100__gpt2/'))

    # args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # _, _, test_loader, (_, _) = get_data_loaders(args, tokenizer, test=True)
    # train_loader, val_loader, dev_val_loader, (train_datasets, test_datasets) = get_data_loaders(args, model.tokenizer)

    # print(f"Loading Model: {args.model_checkpoint}")
    # model.to(args.device)
    # test_model(args,model,tokenizer,test_loader)


if __name__ == "__main__":
    test()