lmflow/pipeline/evaluator.py

"""The Evaluator class simplifies the process of running evaluation on a language model provided by a HFDecoderModel instance imported from the lmflow package. The class constructor takes three dictionaries as arguments: model_args containing arguments related to the language model, data_args containing arguments related to the data used for evaluation, and evaluator_args containing other arguments for the evaluation process.

The class has two methods: create_dataloader() that loads the data from the test file, creates a data loader, and returns it with the size of the data, and evaluate(model) that generates output text given input text. It uses the create_dataloader() method to load the data, iterates over the data in mini-batches, and encodes the input text with the encode() method of the HFDecoderModel class. Then, it generates output text using the evaluate() method of the HFDecoderModel class, decodes the generated output text using the decode() method of the HFDecoderModel class, and writes the output to a file in the output directory. The method also logs some information to the console and Weights and Biases if the use_wandb argument is True.
"""
import os
# import deepspeed
import torch
import wandb
import deepspeed
import sys
import numpy as np
import datetime
import json
# TODO: remove later
from transformers import AutoConfig
import torch.distributed as dist

from lmflow.datasets.dataset import Dataset
from lmflow.pipeline.base_pipeline import BasePipeline
from lmflow.models.hf_decoder_model import HFDecoderModel
from lmflow.utils.data_utils import set_random_seed, batchlize, answer_extraction
os.environ["TOKENIZERS_PARALLELISM"] = "false"  # To avoid warnings about parallelism in tokenizers

class Evaluator(BasePipeline):
    """
    Initializes the `Evaluator` class with given arguments.

    Parameters
    ------------
    model_args : ModelArguments object.
        Contains the arguments required to load the model.
    
    data_args : DatasetArguments object.
        Contains the arguments required to load the dataset.

    evaluator_args : EvaluatorArguments object.
        Contains the arguments required to perform evaluation.


    """
    def __init__(self, model_args, data_args, evaluator_args):
    # our method
        self.data_args = data_args
        self.evaluator_args = evaluator_args
        self.model_args = model_args
        print("--------Begin Evaluator Arguments----------")
        print(f"model_args : {self.model_args}")
        print(f"data_args : {self.data_args}")
        print(f"evaluator_args : {self.evaluator_args}")
        print("--------End Evaluator Arguments----------")
        # logger
        if(self.evaluator_args.use_wandb == True):
            wandb.init(project="lmflow_evaluation")
        # random seed
        set_random_seed(self.evaluator_args.random_seed)
        self.local_rank = int(os.getenv("LOCAL_RANK", "0"))
        self.world_size = int(os.getenv("WORLD_SIZE", "1"))
        torch.cuda.set_device(self.local_rank)  # NOTE: cpu-only machine will have error
        deepspeed.init_distributed()

        self.config = AutoConfig.from_pretrained(model_args.model_name_or_path)
        try: 
            self.model_hidden_size = self.config.hidden_size
        except:
            print("Error in setting hidden size, use the default size 1024")
            self.model_hidden_size = 1024 # gpt2 seems do not have hidden_size in config

        print(f"model_hidden_size = {self.model_hidden_size}")
        # batch size has to be divisible by world_size, but can be bigger than world_size
        train_batch_size = 1 * self.world_size
        self.evaluator_args.minibatch_size = train_batch_size
        self.block_size = evaluator_args.evaluate_block_size
        # dataloader, data_size = create_dataloader(args)    # load dataset


    def create_dataloader(self, dataset: Dataset):
        data_dict = dataset.to_dict()
        inputs = [ instance["input"] for instance in data_dict["instances"] ]
        outputs = [ instance["output"] for instance in data_dict["instances"] ]
        dataset_size = len(outputs)
        dataset_buf = []
        for idx in range(dataset_size):
            dataset_buf.append({
                "input": inputs[idx],
                "output": outputs[idx],
                "input_idx": idx
            })

        dataloader = batchlize(
            dataset_buf,
            self.evaluator_args.minibatch_size,
            self.evaluator_args.random_shuffle
        )
        print(f"Successfully create dataloader with size {len(dataloader)}.")
        return dataloader, dataset_size


    # TODO: Split for better unittest

    def _match(self, predicted_answer, groundtruth, answer_type=None):
        case_insensitive_types = [
            "strategyqa",
            "coin_flip",
            "pubmedqa",
            "binary_choice",
            "medmcqa",
            "usmle",
        ]
        if answer_type in case_insensitive_types:
            return predicted_answer.lower() == groundtruth.lower()
        else:
            return predicted_answer == groundtruth
        return False


    def evaluate(self, model, dataset: Dataset, metric = "accuracy"):
        """
        Perform Evaluation for a model

        Parameters
        ------------
        model : TunableModel object.
            TunableModel to perform inference

        dataset : Dataset object.
            

        """
        if metric in ["acc", "accuracy"]:
            dataloader, data_size = self.create_dataloader(dataset)

            if not dist.is_initialized() or dist.get_rank() == 0:
                if not os.path.exists(self.evaluator_args.output_dir):
                    os.makedirs(self.evaluator_args.output_dir)
                output_writer = open(f"{self.evaluator_args.output_dir}/evaluation.json", "w")

            acc_list = []
            total = 0
            # ds_engine = deepspeed.initialize(model=model.get_model(), config_params=self.ds_config)[0]
            # ds_engine.module.eval()
            for batch_index, batch in enumerate(dataloader):
                if batch_index * self.world_size >= self.data_args.max_eval_samples: 
                    break
                if self.local_rank >= len(batch):
                    current_batch = batch[0]
                else:
                    # the batch in current process
                    current_batch = batch[self.local_rank] 
                
                prompt_structure = self.evaluator_args.prompt_structure
                input = prompt_structure.format(input=current_batch['input'])
                output = current_batch['output']
                input_idx = current_batch['input_idx']

                inputs = model.encode(input, return_tensors="pt").to(device=self.local_rank)
                

                # with torch.no_grad():
                    # outputs = ds_engine.module.generate(inputs, synced_gpus=True, pad_token_id=model.get_tokenizer().eos_token_id, min_length=5, max_length=100,temperature=0.0, do_sample=False)
                outputs = model.inference(inputs, max_new_tokens=100, temperature=0.0)
                text_out = model.decode(outputs[0], skip_special_tokens=True)

                # # only return the generation, trucating the input
                prompt_length = len(model.decode(inputs[0], skip_special_tokens=True,))
                text_out = text_out[prompt_length:]
                answer_type = self.evaluator_args.answer_type
                pred_answer = answer_extraction(
                    text_out,
                    answer_type=answer_type,
                )
                print(f"batch_index{batch_index} rank{self.local_rank}:\n   question={input}\n  prediction={text_out}\n")
                print(f"predicted answer: {pred_answer} \n")
                print(f"groundtruth answer: {output} \n")

                if self.local_rank >= len(batch): # for last batch, the padding examples are ignored and donot contribute to the accuracy
                    correct_ = 0
                    total_ = 0
                else:
                    correct_ = 0
                    total_ = 1
                    if self._match(pred_answer, output, answer_type):
                        correct_ = 1

                # collect accuracy from all gpus
                all_process = torch.tensor([correct_, total_], dtype=torch.float32, device=self.local_rank)
                dist.all_reduce(all_process, dist.ReduceOp.SUM, async_op=False)
                correct_, total_ = all_process.tolist()
                avg = correct_ / total_
                acc_list.append(avg)
                total += total_

                # collect predictions from all gpus
                output_dict = {"question": input,
                            "prediction": text_out,
                            "pred_answer": pred_answer,
                            "answer": output}
                all_process_list = [{}] * self.world_size

                dist.gather_object(output_dict, all_process_list if dist.get_rank() == 0 else None, dst=0)
                if not dist.is_initialized() or dist.get_rank() == 0:
                    current_accuracy = np.mean(acc_list)
                    print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "{}/ {} has been finished, current accuracy = {}".format(int(total), data_size, current_accuracy))
                    
                    if(self.evaluator_args.use_wandb == True):
                        wandb.log({"Accuracy": current_accuracy})

                    for index, output in enumerate(all_process_list):
                        output_json = json.dumps(output)
                        output_writer.write(output_json + '\n')

            if not dist.is_initialized() or dist.get_rank() == 0:
                current_accuracy = np.mean(acc_list)
                print("Final accuracy = ", current_accuracy)
                output_writer.close()
        elif metric in ["ppl", "perplexity"]:
            ppl =  self._evaluate_ppl(model, dataset)
            print(f"Evaluating final ppl: {ppl}")
        elif metric in ["nll", "neg_log_likelihood"]:
            neg_log_likelihood = self._evaluate_neg_log_likelihood(model, dataset)
            print(f"Evaluating final negative log likelihood: {neg_log_likelihood}")
        else:
            raise NotImplementedError(f"{metric} is not implemented or not match with our defined metrics")


    def _evaluate_ppl(self, model, dataset: Dataset):
        data_dict = dataset.to_dict()
        if data_dict['type'] == 'text2text':
            raise NotImplementedError("ppl evaluation is currently not supported for text2text dataset, please use text_only dataset.")
        texts = [ instance["text"] for instance in data_dict["instances"] ]
        encodings = model.get_tokenizer()("\n\n".join(texts), return_tensors="pt")
        # Define some constant
        try:
            max_length = min(model.get_backend_model().config.n_positions, model.get_max_length())
        except:
            max_length = min(1024, model.get_max_length())

        print(f"The maximum sequence length : {max_length}")
        seq_len = encodings.input_ids.size(1)

        nlls = []
        prev_end_loc = 0
        for begin_loc in range(0, seq_len, self.block_size):
            end_loc = min(begin_loc + max_length, seq_len)
            trg_len = end_loc - prev_end_loc  # may be different from block_size on last loop
            input_ids = encodings.input_ids[:, begin_loc:end_loc].to(device=self.local_rank)
            target_ids = input_ids.clone()
            target_ids[:, :-trg_len] = -100

            with torch.no_grad():
                outputs = model.get_backend_model()(input_ids, labels=target_ids)
                # loss is calculated using CrossEntropyLoss which averages over valid labels
                # N.B. the model only calculates loss over trg_len - 1 labels, because it internally shifts the labels
                # to the left by 1.
                neg_log_likelihood = outputs.loss

            nlls.append(neg_log_likelihood)
            prev_end_loc = end_loc
            print(f"Evaluating PPL: {int(begin_loc/self.block_size) + 1} / {int(seq_len/self.block_size)} Complete, current ppl : {torch.exp(torch.stack(nlls).mean())}")
            if end_loc == seq_len:
                break
        ppl = torch.exp(torch.stack(nlls).mean())
        return ppl


    def _evaluate_neg_log_likelihood(self, model, dataset: Dataset):
        """
        Evaluates negative log likelihood of the model over a dataset.

        NLL = -1/N sum_{i=1}^N sum_{j=1}^|w_i| ln(p(w_{i,j}|context_window)),

        where N is the number of data samples, w_{i,j} is the j-th token in
        i-th sample. Here "context_window" = p(w_{i,start}, w_{i,start+1}, ...,
        p_{i,j-1} with start = max(0, j - window_length + 1). "window_length"
        is normally the maximum length accepted by the model.

        Returns:
            A float which represents the negative log likelihood.
        """
        data_dict = dataset.to_dict()

        # Handles prompt structure
        if dataset.get_type() == "text2text":
            prompt = self.evaluator_args.prompt_structure
            data_dict["instances"] = [
                {
                    "input": prompt.format(input=instance["input"]),
                    "output": instance["output"]
                }
                for instance in data_dict["instances"]
            ]

        dataset = dataset.from_dict(data_dict)
        tokenized_dataset = model.tokenize(dataset, add_special_tokens=False)
        tokenized_dataset = tokenized_dataset.get_backend_dataset()
        encoding_list = [
            {
                "input_ids": torch.tensor([input_ids]),
                "labels": torch.tensor([labels]),
            }
            for input_ids, labels in zip(tokenized_dataset["input_ids"],
                                         tokenized_dataset["labels"])
        ]

        # Gets context window length
        try:
            max_length = min(model.get_backend_model().config.n_positions,
                             model.get_max_length())
        except:
            max_length = min(1024, model.get_max_length())

        nlls = []
        full_nlls = []
        num_samples = len(encoding_list)
        for sample_idx, encodings in enumerate(encoding_list):
            seq_len = encodings["input_ids"].size(1)

            prev_end_loc = 0
            for begin_loc in range(0, seq_len, self.block_size):
                end_loc = min(begin_loc + max_length, seq_len)

                # may be different from block_size on last loop
                trg_len = end_loc - prev_end_loc
                input_ids = encodings["input_ids"][:, begin_loc:end_loc]
                input_ids = input_ids.to(device=self.local_rank)

                labels = encodings["labels"][:, begin_loc:end_loc]
                target_ids = labels.clone()
                full_target_ids = input_ids.clone()

                def get_nll(label_ids, nll_list):
                    label_ids[:, :-trg_len] = -100
                    label_ids = label_ids.to(device=self.local_rank)

                    # Valid labels are from 0 to `vocab_size`
                    num_valid_labels = torch.count_nonzero(label_ids >= 0)
                    if label_ids[0, 0] != -100:
                        num_valid_labels -= 1

                    if not torch.all(label_ids == -100):
                        with torch.no_grad():
                            outputs = model.get_backend_model()(
                                input_ids, labels=label_ids
                            )
                            # loss is calculated using CrossEntropyLoss which
                            # sums over valid labels N.B. the model only
                            # calculates loss over trg_len - 1 labels, because
                            # it internally shifts the labels to the left by 1.
                            neg_log_likelihood = outputs.loss * num_valid_labels
                    else:
                        neg_log_likelihood = torch.zeros([]).to(
                            device=self.local_rank
                        )

                    nll_list.append(neg_log_likelihood)

                get_nll(target_ids, nlls)
                get_nll(full_target_ids, full_nlls)

                current_output_nll = torch.stack(nlls).sum() / (sample_idx + 1)
                current_full_nll = torch.stack(full_nlls).sum() / (sample_idx + 1)

                prev_end_loc = end_loc
                if dataset.get_type() == "text_only":
                    print(
                        f"Evaluating negative log likelihood:"
                        f" {sample_idx + 1} / {num_samples} Complete,"
                        f" current nll: {current_full_nll}"
                    )
                elif dataset.get_type() == "text2text":
                    print(
                        f"Evaluating negative log likelihood:"
                        f" {sample_idx + 1} / {num_samples} Complete,"
                        f" current full nll / input nll / output nll:"
                        f" {current_full_nll} /"
                        f" {current_full_nll - current_output_nll} /"
                        f" {current_output_nll}"
                    )
                else:
                    raise NotImplementedError(
                        "f{dataset.get_type()} typed datasets are not supported"
                    )

                if end_loc == seq_len:
                    break

        mean_nll = torch.stack(nlls).sum() / num_samples
        return mean_nll