import yaml import sys import gc import json import os from argparse import Namespace from dataclasses import fields from functools import partial from collections import defaultdict from typing import Callable import numpy as np import torch from rouge_score import rouge_scorer from transformers import ( GenerationConfig, Trainer, Seq2SeqTrainer, Seq2SeqTrainingArguments, EvalPrediction, ) from data_utils import get_tokenized_dataset from modeling_utils import ModulatedPretrainedModel from model_loading import get_tokenizer def clear_gpu(): gc.collect() torch.cuda.empty_cache() torch.cuda.reset_max_memory_allocated() torch.cuda.reset_max_memory_cached() def compute_rouge(pred_texts, label_texts): out = defaultdict(list) scorer = rouge_scorer.RougeScorer(["rouge1", "rougeL"], use_stemmer=False) for pred_text, label_text in zip(pred_texts, label_texts): scores = scorer.score(pred_text, label_text) for k, v in scores.items(): out[f"{k}.f1"].append(v.fmeasure) for k in out: out[k] = np.mean(out[k]) return out def compute_per_token_acc(shift_logits, shift_labels, valid_masks): indices = torch.where(valid_masks) acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item() return {"per_token_acc": acc} def compute_prefix_matching(shift_logits, shift_labels, valid_masks): lengths = valid_masks.sum(dim=1) is_wrong = (shift_logits.argmax(-1) != shift_labels) * valid_masks is_correct = (shift_logits.argmax(-1) == shift_labels) * valid_masks # NOTE: not reliable for multi-turn conversations # ie, all tokens in the following user's turn will be correct # still monotonically correlate with perf though wrong_pos = torch.argmax(is_wrong, dim=1) - torch.argmax(valid_masks, dim=1) perf = wrong_pos / lengths # if all tokens are correct, set to 1 perf = torch.where(is_correct.sum(dim=1) == lengths, 1, perf) return {"prefix_matching": perf.mean().item()} @torch.no_grad() def compute_perplexity(shift_logits, shift_labels, valid_masks): loss_fct = torch.nn.CrossEntropyLoss(reduction="none") loss = loss_fct(shift_logits.transpose(1, 2), shift_labels) loss = (loss * valid_masks).sum(dim=1) / valid_masks.sum(dim=1) perplexity = torch.exp(loss).mean().item() return {"perplexity": perplexity} class Evaluator: def __init__(self, metric_fns: list[Callable]): self.metric_fns = metric_fns self.reset() def reset(self): self.accum_metrics = defaultdict(list) def update(self, shift_logits, shift_labels, valid_masks): for metric_fn in self.metric_fns: metric = metric_fn(shift_logits, shift_labels, valid_masks) for k, v in metric.items(): self.accum_metrics[k].append(v) def compute(self): # Get result across entire eval set result = {k: np.mean(v) for k, v in self.accum_metrics.items()} # Reset batch statistics self.reset() return result @torch.no_grad() def compute_metrics( eval_pred: EvalPrediction, compute_result: bool, evaluator: Evaluator, ): logits, labels = eval_pred.predictions, eval_pred.label_ids shift_logits = logits[..., :-1, :] shift_labels = labels[..., 1:] valid_masks = torch.where(shift_labels != -100, 1, 0) evaluator.update(shift_logits, shift_labels, valid_masks) if compute_result: return evaluator.compute() def save_generated_text(samples, output_dir, split): os.makedirs(output_dir, exist_ok=True) with open(f"{output_dir}/{split}_generated_text.jsonl", "w") as f: for sample in samples: f.write(json.dumps(sample) + "\n") def decode_test_result(test_dataset, test_result, tokenizer): out = [] for sample, pred_toks in zip(test_dataset, test_result.predictions): # sample and labels are not padded # pred_toks are padded though d = dict() if "labels" in sample: start_idx = np.argmax(sample["labels"] != -100) label_toks = sample["labels"][start_idx:] # labels are padded with -100, so we need to # replace them with the pad token id label_toks = np.where(label_toks == -100, tokenizer.pad_token_id, label_toks) label_text = tokenizer.decode(label_toks, skip_special_tokens=True) d["label"] = label_text # remove the label part input_toks = sample["input_ids"][:start_idx] gen_toks = pred_toks[np.argmax(pred_toks != tokenizer.pad_token_id) :] gen_toks = gen_toks[start_idx:] gen_toks = np.where(gen_toks == -100, tokenizer.pad_token_id, gen_toks) d["input"] = tokenizer.decode(input_toks, skip_special_tokens=True) d["generated"] = tokenizer.decode(gen_toks, skip_special_tokens=True) if "ctx_ids" in sample: d["context"] = tokenizer.decode(sample["ctx_ids"], skip_special_tokens=True) out.append(d) return out @torch.no_grad() def eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs): if not isinstance(datasets, dict): datasets = {"": datasets} out = {} for ds_name, ds in datasets.items(): split_name = f"{split}_{ds_name}" if ds_name else split eval_result = eval_trainer.predict( ds, metric_key_prefix=split_name, **gen_kwargs, ) decoded_txts = decode_test_result(ds, eval_result, tokenizer) rouge_metrics = compute_rouge( [txt["generated"] for txt in decoded_txts], [txt["label"] for txt in decoded_txts], ) for k, v in rouge_metrics.items(): eval_result.metrics[f"{split_name}_{k}"] = v save_generated_text( decoded_txts, split=split_name, output_dir=eval_trainer.args.output_dir, ) out[split_name] = eval_result.metrics eval_trainer.log_metrics(split_name, eval_result.metrics) eval_trainer.save_metrics(split_name, eval_result.metrics) clear_gpu() return out @torch.no_grad() def eval_teacher_forcing(eval_trainer, datasets, split): if not isinstance(datasets, dict): datasets = {"": datasets} out = {} for ds_name, ds in datasets.items(): split_name = f"{split}_{ds_name}" if ds_name else split metrics = eval_trainer.evaluate(ds, metric_key_prefix=split_name) out[split_name] = metrics eval_trainer.log_metrics(split_name, metrics) eval_trainer.save_metrics(split_name, metrics) clear_gpu() return out def train_collator(inp_list, tokenizer): # input is a list of tokenized sequences padding_kwargs = dict( padding=True, padding_side="right", pad_to_multiple_of=8, return_tensors="pt", ) ctx_ids = None if "ctx_ids" in inp_list[0]: # have to be manual since it has [ctx_len, features] shape # pad to the longest ctx_len in the batch # which can have a different length from the input_ids, attn_mask, labels ctx_ids = [example.pop("ctx_ids") for example in inp_list] ctx_ids = torch.nn.utils.rnn.pad_sequence( ctx_ids, batch_first=True, padding_value=0, ) # exotic keys won't be padded, so we need to pad them as well ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list] ctx_attn_mask = torch.nn.utils.rnn.pad_sequence( ctx_attn_mask, batch_first=True, padding_value=0, ) chat_ids = None if "chat_ids" in inp_list[0]: chat_ids = [x.pop("chat_ids") for x in inp_list] chat_ids = torch.nn.utils.rnn.pad_sequence( chat_ids, batch_first=True, padding_value=0, ) chat_attn_mask = [x.pop("chat_attn_mask") for x in inp_list] chat_attn_mask = torch.nn.utils.rnn.pad_sequence( chat_attn_mask, batch_first=True, padding_value=0, ) chat_labels = [x.pop("chat_labels") for x in inp_list] chat_labels = torch.nn.utils.rnn.pad_sequence( chat_labels, batch_first=True, padding_value=-100, ) chat_labels = torch.where(chat_attn_mask == 0, -100, chat_labels) labels = [x.pop("labels") for x in inp_list] padded_seq = tokenizer.pad(inp_list, **padding_kwargs) # hacky explicit padding since the labels are not padded by default labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"] labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels) out = {**padded_seq, "labels": labels} if ctx_ids is not None: out["ctx_ids"] = ctx_ids out["ctx_attn_mask"] = ctx_attn_mask if chat_ids is not None: out["chat_ids"] = chat_ids out["chat_attn_mask"] = chat_attn_mask out["chat_labels"] = chat_labels return out def generation_collator(inp_list, tokenizer): padding_kwargs = dict(padding=True, padding_side="left", return_tensors="pt") input_ids = [x.pop("input_ids") for x in inp_list] attn_mask = [x.pop("attention_mask") for x in inp_list] labels = [x.pop("labels") for x in inp_list] for i, label in enumerate(labels): # remove the response tokens idx = np.argmax(label != -100) input_ids[i] = input_ids[i][:idx] attn_mask[i] = attn_mask[i][:idx] out = tokenizer.pad( {"input_ids": input_ids, "attention_mask": attn_mask}, **padding_kwargs ) out["labels"] = labels if "ctx_ids" in inp_list[0]: # have to be manual since it has [ctx_len, features] shape # pad to the longest ctx_len in the batch # which can have a different length from the input_ids, attn_mask, labels ctx_ids = [example.pop("ctx_ids") for example in inp_list] ctx_ids = torch.nn.utils.rnn.pad_sequence( ctx_ids, batch_first=True, padding_value=0, ) # exotic keys won't be padded, so we need to pad them as well ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list] ctx_attn_mask = torch.nn.utils.rnn.pad_sequence( ctx_attn_mask, batch_first=True, padding_value=0, ) out["ctx_ids"] = ctx_ids out["ctx_attn_mask"] = ctx_attn_mask return out def evaluate(checkpoint_path, args, split, generative): assert split in ["validation", "test"] model = ModulatedPretrainedModel.from_state_dict( torch.load(open(checkpoint_path, "rb")), train=False ) tokenizer = get_tokenizer(args.model_name_or_path, train=False) if tokenizer.pad_token_id is None: tokenizer.pad_token_id = tokenizer.eos_token_id model.base_model.config.pad_token_id = tokenizer.pad_token_id model.base_model.generation_config.pad_token_id = tokenizer.pad_token_id ctx_tokenizer = tokenizer if args.ctx_encoder_model_name_or_path: ctx_tokenizer = get_tokenizer(args.ctx_encoder_model_name_or_path, train=False) if ctx_tokenizer.pad_token_id is None: ctx_tokenizer.pad_token_id = ctx_tokenizer.eos_token_id tokenizer_kwargs = {"max_length": args.max_base_len} ctx_tokenizer_kwargs = {"max_length": args.max_ctx_len} add_ctx_to_chat = not isinstance(model, ModulatedPretrainedModel) _get_tokenized_dataset = partial( get_tokenized_dataset, tokenizer=tokenizer, tokenizer_kwargs=tokenizer_kwargs, ctx_tokenizer=ctx_tokenizer, ctx_tokenizer_kwargs=ctx_tokenizer_kwargs, add_ctx_to_chat=add_ctx_to_chat, add_repeat_prompt=False, add_negative_prompt=False, use_kl_loss=False, ) datasets = dict() ds_names = args.val_ds_names if split == "validation" else args.test_ds_names for ds_name in ds_names: datasets[ds_name] = _get_tokenized_dataset(ds_name, split) print(datasets) gen_kwargs = dict(do_sample=False, max_new_tokens=args.max_new_tokens) eval_trainer_args = {} # Copy only necessary attributes from training_args to eval_trainer_args seq2seq_training_args_fields = {f.name for f in fields(Seq2SeqTrainingArguments)} for attr, value in dict(**vars(args)).items(): if attr in seq2seq_training_args_fields and not attr.startswith("_"): eval_trainer_args[attr] = value eval_trainer_args["eval_strategy"] = "no" eval_trainer_args["save_strategy"] = "no" eval_trainer_args["overwrite_output_dir"] = True eval_trainer_args["per_device_eval_batch_size"] = 64 eval_trainer_args = Seq2SeqTrainingArguments( **eval_trainer_args, predict_with_generate=generative, generation_config=GenerationConfig(**gen_kwargs), ) # Seq2SeqTrainer is actually just the same as Trainer # (although it uses a different data collator, i.e., explicit prompt/answer separation) # it just allows `predict_with_generate` # allowing us to compute metrics on the generated outputs # no clue why they call this seq2seq... print("=" * 80 + "\n" + "Evaluating model..." + "\n" + "=" * 80) print(f"checkpoint_path: {checkpoint_path}") model.eval() collator = generation_collator if generative else train_collator trainer_kwargs = { "model": model, "args": eval_trainer_args, "data_collator": partial(collator, tokenizer=tokenizer), } out = {} if not generative: trainer_kwargs["compute_metrics"] = partial( compute_metrics, evaluator=Evaluator( [compute_per_token_acc, compute_prefix_matching, compute_perplexity] ), ) # this is insane # but i don't know why calling trainer.evaluate() on different datasets # always gives the same numbers across datasets... # spents a few hours on this but couldn't find the reason for ds_name, ds in datasets.items(): eval_trainer = Trainer(**trainer_kwargs) metrics = eval_teacher_forcing(eval_trainer, {ds_name: ds}, split) out.update(metrics) else: eval_trainer = Seq2SeqTrainer(**trainer_kwargs) metrics = eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs) out.update(metrics) clear_gpu() return out if __name__ == "__main__": os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" os.environ["WANDB_MODE"] = "disabled" checkpoint_path = sys.argv[1] checkpoint_dir = "/".join(checkpoint_path.split("/")[:-1]) run_dir = "/".join(checkpoint_path.split("/")[:-2]) cur_it = int(checkpoint_path.split("checkpoint-")[1].split("/")[0]) args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml", "r"))) print(f"checkpoint_path: {checkpoint_path}") print(f"run_dir: {run_dir}") # print(f"args: {args}") args.output_dir = f"{run_dir}/eval-results-{cur_it}" args.logging_dir = f"{run_dir}/eval-results-{cur_it}" args.run_name = run_dir.split("/")[-1] evaluate(checkpoint_path, args, split="validation", generative=False) evaluate(checkpoint_path, args, split="validation", generative=True)