Source code for lmflow.pipeline.evaluator

"""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 torch
import wandb
import deepspeed
import sys
import numpy as np
import datetime
import json
# TODO: remove later
from accelerate import Accelerator
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

[docs] 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
[docs] self.data_args = data_args
[docs] self.evaluator_args = evaluator_args
[docs] self.model_args = model_args
# logger if(self.evaluator_args.use_wandb == True): wandb.init(project="lmflow_evaluation") # random seed set_random_seed(self.evaluator_args.random_seed)
[docs] self.local_rank = int(os.getenv("LOCAL_RANK", "0"))
[docs] self.world_size = int(os.getenv("WORLD_SIZE", "1"))
torch.cuda.set_device(self.local_rank) # NOTE: cpu-only machine will have error if evaluator_args.use_accelerator_for_evaluator: self.accelerator = Accelerator() self.accelerator.wait_for_everyone() else: deepspeed.init_distributed()
[docs] 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 = self.evaluator_args.inference_batch_size_per_device * self.world_size self.evaluator_args.minibatch_size = train_batch_size
[docs] self.block_size = evaluator_args.evaluate_block_size
# dataloader, data_size = create_dataloader(args) # load dataset
[docs] 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)},batch_size {self.evaluator_args.minibatch_size}.") return dataloader, dataset_size
# TODO: Split for better unittest
[docs] 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
[docs] def evaluate( self, model, dataset: Dataset, metric = "accuracy", verbose=True, ): """ Perform Evaluation for a model Parameters ------------ model : TunableModel object. TunableModel to perform inference dataset : Dataset object. """ if metric in ["acc", "accuracy"]: if self.evaluator_args.use_accelerator_for_evaluator: acc = self._evaluate_acc_with_accelerator(model, dataset, verbose=verbose) else: acc = self._evaluate_acc_with_deepspeed(model, dataset, verbose=verbose) print(f"Evaluating final accuracy: {acc}") return acc elif metric in ["ppl", "perplexity"]: ppl = self._evaluate_ppl(model, dataset, verbose=verbose) print(f"Evaluating final perplexity: {ppl}") return ppl elif metric in ["nll", "neg_log_likelihood"]: nll = self._evaluate_nll(model, dataset, verbose=verbose) print(f"Evaluating final negative log likelihood: {nll}") return nll else: raise NotImplementedError(f"metric {metric} is not supported")
[docs] def _evaluate_acc_with_accelerator(self, model, dataset, verbose=True): dataloader, data_size = self.create_dataloader(dataset) if self.accelerator.is_local_main_process: 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") correct_number_list = [] for batch_index, batch in enumerate(dataloader): if batch_index * self.world_size >= self.data_args.max_eval_samples: break if self.local_rank*self.evaluator_args.inference_batch_size_per_device >= len(batch): current_batch = batch[:self.evaluator_args.inference_batch_size_per_device] else: current_batch = batch[self.local_rank*self.evaluator_args.inference_batch_size_per_device:(self.local_rank+1)*self.evaluator_args.inference_batch_size_per_device] prompt_structure = self.evaluator_args.prompt_structure input = [prompt_structure.format(input=i['input']) for i in current_batch] output = [i['output'] for i in current_batch] batch_input = model.encode(input, return_tensors="pt",padding=True).to(self.accelerator.device) inputs = batch_input['input_ids'] mask = batch_input['attention_mask'] with self.accelerator.autocast(): outputs = model.inference(inputs, max_new_tokens=self.evaluator_args.max_new_tokens,attention_mask=mask,temperature=self.evaluator_args.temperature, repetition_penalty=self.evaluator_args.repetition_penalty,use_accelerator=self.evaluator_args.use_accelerator_for_evaluator) text_out = model.decode(outputs, skip_special_tokens=True) decoded_input = model.decode(inputs, skip_special_tokens=True,) prompt_length = [len(i) for i in decoded_input] text_out = [text_out[i][prompt_length[i]:] for i in range(len(text_out))] answer_type = self.evaluator_args.answer_type pred_answer = [] for i in text_out: pred_answer.append(answer_extraction( i, answer_type=answer_type, )) if verbose: 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 * self.evaluator_args.inference_batch_size_per_device >= len(batch): correct_ = 0 else: correct_ = 0 for i in range(len(pred_answer)): if self._match(pred_answer[i], output[i], answer_type): correct_ += 1 # collect accuracy from all gpus all_process = torch.tensor([correct_], dtype=torch.float32, device=self.local_rank) all_process = self.accelerator.gather(all_process) correct_ = sum(all_process.tolist()) correct_number_list.append(correct_) # collect predictions from all gpus output_dict = {"question": input, "prediction": text_out, "pred_answer": pred_answer, "answer": output} if(self.world_size > 1): all_process_list = [{}] * self.world_size dist.gather_object(output_dict, all_process_list if dist.get_rank() == 0 else None, dst=0) else: all_process_list = [output_dict] if self.accelerator.is_local_main_process: current_total = (batch_index+1) * self.world_size * self.evaluator_args.inference_batch_size_per_device current_accuracy = np.sum(correct_number_list) / current_total if int(current_total) < data_size else np.sum(correct_number_list) / data_size print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), f"{int(current_total) if int(current_total) < data_size else data_size} / {data_size} has been finished, # correct = { np.sum(correct_number_list)}, current accuracy = {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 self.accelerator.is_local_main_process: current_accuracy = np.sum(correct_number_list) / data_size print(f"# Correct = {np.sum(correct_number_list)}, # Total = {data_size}, Final accuracy = ", current_accuracy) output_writer.close() return np.sum(correct_number_list) / data_size
[docs] def _evaluate_acc_with_deepspeed(self, model, dataset, verbose=True): 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") correct_number_list = [] for batch_index, batch in enumerate(dataloader): if batch_index * self.world_size >= self.data_args.max_eval_samples: break if self.local_rank*self.evaluator_args.inference_batch_size_per_device >= len(batch): current_batch = batch[:self.evaluator_args.inference_batch_size_per_device] else: current_batch = batch[self.local_rank*self.evaluator_args.inference_batch_size_per_device:(self.local_rank+1)*self.evaluator_args.inference_batch_size_per_device] prompt_structure = self.evaluator_args.prompt_structure input = [prompt_structure.format(input=i['input']) for i in current_batch] output = [i['output'] for i in current_batch] input_idx = [i['input_idx'] for i in current_batch] batch_input = model.encode(input, return_tensors="pt",padding=True).to(device=self.local_rank) inputs = batch_input['input_ids'] mask = batch_input['attention_mask'] outputs = model.inference(inputs, max_new_tokens=self.evaluator_args.max_new_tokens, attention_mask=mask,temperature=self.evaluator_args.temperature, repetition_penalty=self.evaluator_args.repetition_penalty) text_out = model.decode(outputs, skip_special_tokens=True) # # only return the generation, trucating the input decoded_input = model.decode(inputs, skip_special_tokens=True,) prompt_length = [len(i) for i in decoded_input] text_out = [text_out[i][prompt_length[i]:] for i in range(len(text_out))] answer_type = self.evaluator_args.answer_type pred_answer = [] for i in text_out: pred_answer.append(answer_extraction( i, answer_type=answer_type, )) if verbose: 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 * self.evaluator_args.inference_batch_size_per_device >= len(batch): correct_ = 0 else: correct_ = 0 for i in range(len(pred_answer)): if self._match(pred_answer[i], output[i], answer_type): correct_ += 1 # collect accuracy from all gpus all_process = torch.tensor([correct_], dtype=torch.float32, device=self.local_rank) dist.all_reduce(all_process, dist.ReduceOp.SUM, async_op=False) correct_ = all_process.tolist() correct_number_list.append(correct_) # 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_total = (batch_index+1) * self.world_size * self.evaluator_args.inference_batch_size_per_device current_accuracy = np.sum(correct_number_list) / current_total if int(current_total) < data_size else np.sum(correct_number_list) / data_size print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), f"{int(current_total) if int(current_total) < data_size else data_size} / {data_size} has been finished, # correct = { np.sum(correct_number_list)}, current accuracy = {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.sum(correct_number_list) / data_size print(f"# Correct = {np.sum(correct_number_list)}, # Total = {data_size}, Final accuracy = ", current_accuracy) output_writer.close() return np.sum(correct_number_list) / data_size
[docs] def _evaluate_ppl(self, model, dataset: Dataset, verbose=True): 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 if self.model_args.truncate_to_model_max_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()) else: max_length = self.block_size if verbose: 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 if verbose: 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
[docs] def _evaluate_nll( self, model, dataset: Dataset, verbose=True, ): """ 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 verbose: 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