mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
427 lines
15 KiB
Python
427 lines
15 KiB
Python
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)
|