From 40a92f482a8facf9c4721f91ad10886e9fb4f262 Mon Sep 17 00:00:00 2001 From: 51616 Date: Tue, 14 Jan 2025 17:47:04 +0000 Subject: [PATCH] load_state_dict correctly + eval script (crazy trainer eval bug...) --- hyperlora/eval.py | 416 ++++++++++++++++++++++++++++++++++++ hyperlora/intx_sft.py | 65 ++---- hyperlora/modeling_utils.py | 20 +- hyperlora/training_utils.py | 80 +------ 4 files changed, 445 insertions(+), 136 deletions(-) create mode 100644 hyperlora/eval.py diff --git a/hyperlora/eval.py b/hyperlora/eval.py new file mode 100644 index 0000000..4f22195 --- /dev/null +++ b/hyperlora/eval.py @@ -0,0 +1,416 @@ +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): + 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} + + 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, + ) + eval_trainer.log_metrics(split_name, eval_result.metrics) + eval_trainer.save_metrics(split_name, eval_result.metrics) + clear_gpu() + + +@torch.no_grad() +def eval_teacher_forcing(eval_trainer, datasets, split): + if not isinstance(datasets, dict): + datasets = {"": datasets} + + 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) + eval_trainer.log_metrics(split_name, metrics) + eval_trainer.save_metrics(split_name, metrics) + clear_gpu() + + +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), + } + 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) + eval_teacher_forcing(eval_trainer, {ds_name: ds}, split) + else: + eval_trainer = Seq2SeqTrainer(**trainer_kwargs) + eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs) + + clear_gpu() + + +if __name__ == "__main__": + os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" + checkpoint_path = sys.argv[1] + checkpoint_dir = "/".join(checkpoint_path.split("/")[:-1]) + run_dir = "/".join(checkpoint_path.split("/")[:-2]) + 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 = checkpoint_dir + args.logging_dir = checkpoint_dir + args.run_name = run_dir.split("/")[-1] + + evaluate(checkpoint_path, args, split="validation", generative=False) + evaluate(checkpoint_path, args, split="validation", generative=True) diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index ea9dee5..72e75bf 100755 --- a/hyperlora/intx_sft.py +++ b/hyperlora/intx_sft.py @@ -80,7 +80,7 @@ LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) 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} + return {"perplexity": acc} def compute_prefix_matching(shift_logits, shift_labels, valid_masks): @@ -222,18 +222,6 @@ def main(): ] ) - args = { - **vars(deepcopy(data_args)), - **vars(deepcopy(ctx_args)), - **vars(deepcopy(model_args)), - **vars(deepcopy(lora_args)), - **vars(deepcopy(training_args)), - **vars(deepcopy(hypernet_args)), - **vars(deepcopy(aggregator_args)), - **vars(deepcopy(ctx_encoder_args)), - } - args["deepspeed_plugin"] = None - checkpoint_dir = training_args.resume_from_checkpoint if checkpoint_dir and not os.path.isdir(checkpoint_dir): raise NotADirectoryError(f"Checkpoint{checkpoint_dir} is not a directory") @@ -242,12 +230,12 @@ def main(): # still possible to have a name crash though # logging_dir is just "runs/DATE_TIME_HOSTNAME" run_name = ( - get_run_name(seed_str=training_args.logging_dir) + get_run_name(seed_str=training_args.logging_dir.strip("runs/")) if not checkpoint_dir else checkpoint_dir.split("/")[-2] ) - output_dir = f"train_outputs/{run_name}" + output_dir = f"train_outputs/runs/{run_name}" setup_logging(output_dir, debug=os.getenv("DEBUG", False)) logger.debug(f"CMD: {' '.join(os.sys.argv)}") cli_args = extract_cli_args(os.sys.argv) @@ -266,6 +254,17 @@ def main(): and training_args.lr_scheduler_kwargs is None ): training_args.lr_scheduler_kwargs = {"min_lr": 1e-7} + args = { + **vars(deepcopy(data_args)), + **vars(deepcopy(ctx_args)), + **vars(deepcopy(model_args)), + **vars(deepcopy(lora_args)), + **vars(deepcopy(training_args)), + **vars(deepcopy(hypernet_args)), + **vars(deepcopy(aggregator_args)), + **vars(deepcopy(ctx_encoder_args)), + } + args["deepspeed_plugin"] = None logger.debug(f"args: {args}") save_yaml(args, f"{output_dir}/args.yaml") set_seed(training_args.seed) @@ -455,42 +454,6 @@ def main(): 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 - # TODO: use SFTTrainer instead? https://huggingface.co/docs/trl/en/sft_trainer # TODO: use packing with SFTTrainer diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py index c6ce671..be2f4fa 100644 --- a/hyperlora/modeling_utils.py +++ b/hyperlora/modeling_utils.py @@ -575,9 +575,11 @@ class ModulatedPretrainedModel(nn.Module): requires_grad=False, peft_config=lora_config, ) - hypernet_config = state_dict.pop("hypernet_config") - ctx_encoder_args = state_dict.pop("ctx_encoder_args") - return cls(base_model, hypernet_config, ctx_encoder_args) + hypernet_config = state_dict["hypernet_config"] + ctx_encoder_args = state_dict["ctx_encoder_args"] + model = cls(base_model, hypernet_config, ctx_encoder_args) + model.load_state_dict(state_dict) + return model def _init_model(self): self.hypernet = HyperLoRA(self.hypernet_config).to(self.device) @@ -862,7 +864,7 @@ if __name__ == "__main__": # set torch randomness seed torch.manual_seed(42) - model_name = "meta-llama/Llama-3.2-1B-Instruct" + model_name = "meta-llama/Llama-3.1-8B-Instruct" base_model, tokenizer = get_model_and_tokenizer( model_name, train=True, @@ -871,13 +873,16 @@ if __name__ == "__main__": ) # TODO: base model shoul be init with target lora config print(base_model) + ctx_model = base_model device = base_model.device # lora_config = base_model.peft_config["default"] # d_in, d_out = get_peft_in_out_features(base_model, peft_config=lora_config) hypernet_args = HypernetArguments(latent_size=512) aggregator_args = AggregatorArguments(aggregator_type=AGGREGATOR_TYPE.PERCEIVER) - hypernet_config = get_hypernet_config(base_model, hypernet_args, aggregator_args) + hypernet_config = get_hypernet_config( + base_model, ctx_model, hypernet_args, aggregator_args + ) ctx_encoder_args = CtxEncoderArguments(layer_idx=4) # ctx_encoder = EarlyExit(get_base_model(base_model), 4) # hypernet = HyperLoRA( @@ -919,7 +924,10 @@ if __name__ == "__main__": # model.load_state_dict( # torch.load( - # open("train_outputs/20250108-193257_BB1p30QQ/pytorch_model.bin", "rb") + # open( + # "train_outputs/runs/Jan14_11-10-56_slurm0-a3nodeset-1_30c9e93d/checkpoint-1720/pytorch_model.bin", + # "rb", + # ) # ) # ) # modelout = model(ctx_ids, ctx_attn_mask, **prompt_inputs) diff --git a/hyperlora/training_utils.py b/hyperlora/training_utils.py index b6a67ba..35dd398 100644 --- a/hyperlora/training_utils.py +++ b/hyperlora/training_utils.py @@ -28,84 +28,6 @@ def clear_gpu(): 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 save_generated_text(samples, output_dir, split): - 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): - 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[len(input_toks) :] - 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 - - -def eval_generation(eval_trainer, tokenizer, dataset, split, gen_kwargs): - if not isinstance(dataset, dict): - dataset = {"": dataset} - - for ds_name, ds in dataset.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}_{k}"] = v - - save_generated_text( - decoded_txts, - split=split_name, - output_dir=eval_trainer.args.output_dir, - ) - eval_trainer.log_metrics(split_name, eval_result.metrics) - eval_trainer.save_metrics(split_name, eval_result.metrics) - clear_gpu() - - -# def per_sample_loss_avg_fn(outputs, labels, num_items_in_batch): -# ... - - def train_model( model, # tokenizer, @@ -138,8 +60,8 @@ def train_model( # Trainer loads the best model after training # is done when load_best_model_at_end=True train_result = trainer.train(resume_from_checkpoint=checkpoint) - trainer.log_metrics("train", train_result.metrics) trainer.save_model() # just in case OOM when run trainer.evaluate() + trainer.log_metrics("train", train_result.metrics) clear_gpu() metrics = trainer.evaluate(dict(**val_dataset, test=test_dataset)) trainer.log_metrics("eval", metrics)