import gc import json import logging import os import re import string from argparse import Namespace from collections import Counter, defaultdict from collections.abc import Callable from dataclasses import fields from functools import partial import numpy as np import torch import yaml from datasets import disable_caching from peft import PeftModel from rouge_score import rouge_scorer from transformers import ( EvalPrediction, GenerationConfig, PreTrainedModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, Trainer, set_seed, ) from transformers.utils import is_liger_kernel_available from ctx_to_lora.data.definitions import LONGBENCH_E_TASKS, LONGBENCH_TASKS from ctx_to_lora.data.processing import get_tokenized_dataset from ctx_to_lora.model_loading import get_model, get_tokenizer from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel logger = logging.getLogger() # TODO: add negative samples eval (per-sample?) # TODO: add length info for more fine grain eval # longbench metrics def normalize_answer(s): """Lower text and remove punctuation, articles and extra whitespace.""" def remove_articles(text): return re.sub(r"\b(a|an|the)\b", " ", text) def white_space_fix(text): return " ".join(text.split()) def remove_punc(text): exclude = set(string.punctuation) return "".join(ch for ch in text if ch not in exclude) def lower(text): return text.lower() return white_space_fix(remove_articles(remove_punc(lower(s)))) def f1_score(prediction: str, ground_truth: str): common = Counter(prediction) & Counter(ground_truth) num_same = sum(common.values()) if num_same == 0: return 0 precision = 1.0 * num_same / len(prediction) recall = 1.0 * num_same / len(ground_truth) f1 = (2 * precision * recall) / (precision + recall) return f1 def compute_qa_f1_score(pred_texts: list[str], label_texts: list[str]): """ Word-level F1 score for evaluating question answering systems. Order of the words does not matter. """ res = [] for prediction, label in zip(pred_texts, label_texts): normalized_prediction = normalize_answer(prediction) normalized_label = normalize_answer(label) prediction_words = normalized_prediction.split() label_words = normalized_label.split() res.append(f1_score(prediction_words, label_words)) return dict(qa_f1=np.mean(res)) CLOSED_QA_DATASETS = { "longbench/narrativeqa", "longbench/qasper", "longbench/multifieldqa_en", "longbench/hotpotqa", "longbench/2wikimqa", "longbench/musique", "hotpot_qa", "squad", "triviaqa_retrieved", "negative_nq", } for ds_name in list(CLOSED_QA_DATASETS): CLOSED_QA_DATASETS.add(f"{ds_name}_e") def clear_gpu(): gc.collect() torch.cuda.empty_cache() torch.cuda.reset_max_memory_allocated() torch.cuda.reset_max_memory_cached() def add_longbench_tasks(ds_names: list[str]): if "longbench" in ds_names: ds_names.remove("longbench") ds_names += LONGBENCH_TASKS if "longbench_e" in ds_names: ds_names.remove("longbench_e") ds_names += LONGBENCH_E_TASKS 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} acc = (shift_logits.argmax(-1) == shift_labels)[indices].float() return { "per_token_accs": acc.flatten().tolist(), "n_per_token_accs": valid_masks.sum().item(), } 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()} return { "prefix_matchings": perf.tolist(), "n_prefix_matchings": valid_masks.shape[0], } @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} preplexities = torch.exp(loss) return { "perplexities": preplexities.tolist(), "n_perplexities": valid_masks.shape[0], } class Evaluator: def __init__(self, metric_fns: list[Callable]): self.metric_fns = metric_fns self.reset() def reset(self): self.accum_metrics = defaultdict(list) self.count = 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(): if k.startswith("n_"): self.count[k[2:]].append(v) else: self.accum_metrics[k] += v def compute(self): # Get result across entire eval set result = { k: np.sum(v) / np.sum(self.count[k]) for k, v in self.accum_metrics.items() } # Reset batch statistics self.reset() return result @torch.inference_mode() 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) # Create any necessary subdirectories if split contains path separators if "/" in split: split_dir = os.path.join(output_dir, os.path.dirname(split)) os.makedirs(split_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, ctx_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.strip() # 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).strip() if "ctx_ids" in sample: d["context"] = ctx_tokenizer.decode( sample["ctx_ids"], skip_special_tokens=True ) for k in sample: if k.endswith("_len"): d[k] = sample[k].item() out.append(d) # sort samples by length if possible len_key = "ctx_ids_len" if "ctx_ids_len" in sample else "input_ids_len" sorted(out, key=lambda x: x[len_key]) return out @torch.inference_mode() def eval_generation( eval_trainer, tokenizer, ctx_tokenizer, datasets, split, remove_context, gen_kwargs ): if not isinstance(datasets, dict): datasets = {"": datasets} out = {} for ds_name, ds in datasets.items(): print(f"Evaluating: {ds_name}") split_name = f"{split}_{ds_name}" if ds_name else split if remove_context: split_name += "_no_context" eval_result = eval_trainer.predict( ds, metric_key_prefix=split_name, **gen_kwargs, ) decoded_txts = decode_test_result(ds, eval_result, tokenizer, ctx_tokenizer) pred_texts = [txt["generated"] for txt in decoded_txts] label_texts = [txt["label"] for txt in decoded_txts] rouge_metrics = compute_rouge(pred_texts, label_texts) for k, v in rouge_metrics.items(): eval_result.metrics[f"{split_name}_{k}"] = v if ds_name in CLOSED_QA_DATASETS: print("Computing QA F1 Score") qa_f1_metric = compute_qa_f1_score( [txt["generated"] for txt in decoded_txts], [txt["label"] for txt in decoded_txts], ) for k, v in qa_f1_metric.items(): eval_result.metrics[f"{split_name}_{k}"] = v # Group by input length and compute metrics for each group length_bins = [(0, 511), (512, 1023), (1024, 2047), (2048, 4095), (4096, 8192)] # Ensure all keys for length metrics are present, even if a bin is empty for low, high in length_bins: eval_result.metrics[f"{split_name}_rouge1.f1_len_{low}-{high}"] = "None" eval_result.metrics[f"{split_name}_rougeL.f1_len_{low}-{high}"] = "None" if ds_name in CLOSED_QA_DATASETS: eval_result.metrics[f"{split_name}_qa_f1_len_{low}-{high}"] = "None" grouped_texts = defaultdict(lambda: {"generated": [], "label": [], "count": 0}) for txt in decoded_txts: len_key = "ctx_ids_len" if "ctx_ids_len" in txt else "input_ids_len" input_len = txt[len_key] for low, high in length_bins: if low <= input_len <= high: group_key = f"{low}-{high}" grouped_texts[group_key]["generated"].append(txt["generated"]) grouped_texts[group_key]["label"].append(txt["label"]) grouped_texts[group_key]["count"] += 1 break for group_key, data in grouped_texts.items(): if data["count"] > 0: group_rouge_metrics = compute_rouge(data["generated"], data["label"]) for k, v in group_rouge_metrics.items(): eval_result.metrics[f"{split_name}_{k}_len_{group_key}"] = v if ds_name in CLOSED_QA_DATASETS: group_qa_f1_metric = compute_qa_f1_score( data["generated"], data["label"] ) for k, v in group_qa_f1_metric.items(): eval_result.metrics[f"{split_name}_{k}_len_{group_key}"] = 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, remove_context): 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 if remove_context: split_name += "_no_context" 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 ) # we don't include the labels in the output # since we don't want to compute the loss on the labels # during generation 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, model_name_or_path, eval_batch_size, args, split, generative, ): assert split in ["validation", "test"] ctx_name = None if model_name_or_path is None: state_dict = torch.load(checkpoint_path, weights_only=False) ctx_name = state_dict["ctx_encoder_args"].ctx_encoder_model_name_or_path model = ModulatedPretrainedModel.from_state_dict( state_dict, train=False, use_flash_attn=True, ) # if generative: # model = model.to(torch.bfloat16) else: model_kwargs = dict(attn_implementation="flash_attention_2") model = get_model( model_name_or_path, train=False, requires_grad=False, model_kwargs=model_kwargs, ) # NOTE: there is still some randomness in the eval result # despite all the deterministic settings if is_liger_kernel_available(): from liger_kernel.transformers import _apply_liger_kernel_to_instance if isinstance(model, ModulatedPretrainedModel): print("Applying liger-kernel to ModulatedPretrainedModel") if isinstance(model.base_model, PeftModel): _apply_liger_kernel_to_instance(model=model.base_model.base_model.model) else: _apply_liger_kernel_to_instance(model=model.base_model.model) if ctx_name is not None: print("Applying liger-kernel to ctx_encoder_model") _apply_liger_kernel_to_instance(model=model.ctx_encoder.base_model) elif isinstance(model, PeftModel): print("Applying liger-kernel to PeftModel") _apply_liger_kernel_to_instance(model=model.base_model.model) elif isinstance(model, PreTrainedModel): print("Applying liger-kernel to PretrainedModel") _apply_liger_kernel_to_instance(model=model.model) 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 base_model = ( model.base_model if isinstance(model, ModulatedPretrainedModel) else model ) base_model.config.pad_token_id = tokenizer.pad_token_id base_model.generation_config.pad_token_id = tokenizer.pad_token_id ctx_tokenizer = tokenizer if ctx_name: ctx_tokenizer = get_tokenizer(ctx_name, train=False) if ctx_tokenizer.pad_token_id is None: ctx_tokenizer.pad_token_id = ctx_tokenizer.eos_token_id tokenizer_kwargs = {} ctx_tokenizer_kwargs = {} add_ctx_to_chat = ( not isinstance(model, ModulatedPretrainedModel) and not args.remove_context ) ctx_model_max_len = ( model.ctx_encoder.config.max_position_embeddings if isinstance(model, ModulatedPretrainedModel) else None ) _get_tokenized_dataset = partial( get_tokenized_dataset, base_model_max_len=model.base_model.config.max_position_embeddings, tokenizer=tokenizer, tokenizer_kwargs=tokenizer_kwargs, ctx_model_max_len=ctx_model_max_len, 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, set_format="pt", ) datasets = dict() ds_names = args.val_ds_names if split == "validation" else args.test_ds_names add_longbench_tasks(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=256, # 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["batch_eval_metrics"] = True eval_trainer_args["per_device_eval_batch_size"] = eval_batch_size 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, args.remove_context ) out.update(metrics) else: eval_trainer = Seq2SeqTrainer(**trainer_kwargs) metrics = eval_generation( eval_trainer, tokenizer, ctx_tokenizer, datasets, split, args.remove_context, gen_kwargs, ) out.update(metrics) clear_gpu() return out if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Evaluate a checkpoint") parser.add_argument( "--model_name_or_path", type=str, default=None, help="Evaluate a base model from HuggingFace Hub, without loading checkpoint", ) parser.add_argument( "--checkpoint_path", type=str, default=None, help="Path to the checkpoint to evaluate", ) parser.add_argument( "--split", type=str, choices=["validation", "test"], default="validation", help="Which split to evaluate on", ) parser.add_argument( "--datasets", type=str, nargs="+", help=( "Specific datasets to evaluate on." "If not provided, uses default from args.yaml" ), ) parser.add_argument( "--eval_batch_size", type=int, default=8, help="Eval batch size for teacher forcing", ) parser.add_argument( "--eval_batch_size_gen", type=int, default=32, help="Eval batch size for generation", ) parser.add_argument( "--remove_context", action="store_true", help="Remove context when evaluating the base model.", ) cli_args = parser.parse_args() # setup_logging(output_dir, debug=os.getenv("DEBUG", False)) assert bool(cli_args.model_name_or_path) ^ bool(cli_args.checkpoint_path), ( "Either --model_name_or_path or --checkpoint_path must be provided" ) os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1" os.environ["WANDB_MODE"] = "disabled" disable_caching() set_seed(42) torch.use_deterministic_algorithms(True, warn_only=True) torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False torch.backends.cudnn.benchmark = False torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False if cli_args.checkpoint_path: checkpoint_dir = "/".join(cli_args.checkpoint_path.split("/")[:-1]) run_dir = "/".join(cli_args.checkpoint_path.split("/")[:-2]) cur_it = int(cli_args.checkpoint_path.split("checkpoint-")[1].split("/")[0]) args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml"))) print(f"checkpoint_path: {cli_args.checkpoint_path}") print(f"run_dir: {run_dir}") 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] args.remove_context = False # modulated model doesn't see ctx by default else: args = Namespace( model_name_or_path=cli_args.model_name_or_path, output_dir=f"eval_results/{cli_args.model_name_or_path}", logging_dir=f"eval_results/{cli_args.model_name_or_path}", run_name=f"eval_results/{cli_args.model_name_or_path}", # max_base_len=2**13, # max_ctx_len=-1, # not used val_ds_names=[], test_ds_names=[], remove_context=cli_args.remove_context, ) # Override dataset names if provided via CLI if cli_args.datasets: if cli_args.split == "validation": args.val_ds_names = cli_args.datasets else: args.test_ds_names = cli_args.datasets # evaluate( # cli_args.checkpoint_path, # cli_args.model_name_or_path, # cli_args.eval_batch_size, # args, # split=cli_args.split, # generative=False, # ) evaluate( cli_args.checkpoint_path, cli_args.model_name_or_path, cli_args.eval_batch_size_gen, args, split=cli_args.split, generative=True, )