mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
955 lines
33 KiB
Python
955 lines
33 KiB
Python
import gc
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import string
|
|
import sys
|
|
from argparse import Namespace
|
|
from collections import Counter, defaultdict
|
|
from collections.abc import Callable
|
|
from dataclasses import fields
|
|
from functools import partial
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
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 import hypernet
|
|
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
|
|
|
|
# bandaid for loading old models (before restructure)
|
|
sys.modules["ctx_to_lora.modeling_utils"] = hypernet
|
|
|
|
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(["rougeL"], use_stemmer=True)
|
|
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 create_metrics_csv(
|
|
metrics_dict: dict[str, dict[str, Any]],
|
|
output_dir: str,
|
|
model_name: str,
|
|
is_hypernet_model: bool = False,
|
|
remove_context: bool = False,
|
|
) -> None:
|
|
"""
|
|
Create a human-readable CSV file from evaluation metrics with hierarchical columns.
|
|
One row per model with columns structured as task_metric_lengthgroup.
|
|
|
|
Args:
|
|
metrics_dict: Dictionary containing evaluation metrics for each dataset/split
|
|
output_dir: Directory to save the CSV file
|
|
model_name: Name of the model being evaluated
|
|
is_hypernet_model: Whether this is a hypernet/modulated model or base model
|
|
remove_context: Whether context was removed during evaluation
|
|
"""
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# Collect all unique metric names, length groups, and tasks
|
|
all_metrics = set()
|
|
all_length_groups = set()
|
|
all_tasks = set()
|
|
|
|
# FIXME: fix name split (currently wrong)
|
|
for split_name, metrics in metrics_dict.items():
|
|
task_name = split_name.split("_")[1] if "_" in split_name else split_name
|
|
all_tasks.add(task_name)
|
|
|
|
for metric_key in metrics.keys():
|
|
if not metric_key.startswith(split_name):
|
|
continue
|
|
|
|
# Skip timing and performance metrics that aren't evaluation results
|
|
if (
|
|
"model_preparation_time" in metric_key
|
|
or "steps_per_second" in metric_key
|
|
or "samples_per_second" in metric_key
|
|
or "runtime" in metric_key
|
|
):
|
|
continue
|
|
|
|
# Remove the split prefix to get the actual metric name
|
|
metric_name = metric_key[len(split_name) + 1 :]
|
|
|
|
# Check if this is a length-specific metric
|
|
if "_len_" in metric_name:
|
|
base_metric, length_part = metric_name.split("_len_")
|
|
all_metrics.add(base_metric)
|
|
all_length_groups.add(length_part)
|
|
else:
|
|
all_metrics.add(metric_name)
|
|
|
|
# Add overall metric (no length grouping)
|
|
all_length_groups.add("overall")
|
|
|
|
# Sort for consistent ordering
|
|
all_metrics = sorted(all_metrics)
|
|
|
|
# Custom sorting for length groups to ensure proper numerical ordering
|
|
def sort_length_groups(length_group: str) -> tuple:
|
|
if length_group == "overall":
|
|
return (1, 0, 0) # Put "overall" after numerical ranges
|
|
try:
|
|
# Parse "low-high" format
|
|
low, high = map(int, length_group.split("-"))
|
|
return (0, low, high)
|
|
except (ValueError, IndexError):
|
|
return (2, 0, 0) # Put any malformed strings last
|
|
|
|
all_length_groups = sorted(list(all_length_groups), key=sort_length_groups)
|
|
all_tasks = sorted(all_tasks)
|
|
|
|
# Create single row with hierarchical columns
|
|
row_data = {
|
|
"model": model_name,
|
|
"model_type": "hypernet" if is_hypernet_model else "base",
|
|
}
|
|
|
|
# Create hierarchical column structure: task_metric_lengthgroup
|
|
for task in all_tasks:
|
|
# Find the corresponding split in metrics_dict
|
|
split_key = None
|
|
for split_name in metrics_dict.keys():
|
|
if task in split_name:
|
|
split_key = split_name
|
|
break
|
|
|
|
if split_key is None:
|
|
continue
|
|
|
|
metrics = metrics_dict[split_key]
|
|
|
|
# Determine if this task uses QA F1 or ROUGE metrics
|
|
is_qa_task = task in CLOSED_QA_DATASETS
|
|
|
|
for metric in all_metrics:
|
|
# Skip inappropriate metrics for the task type
|
|
if is_qa_task and metric.startswith("rougeL"):
|
|
continue
|
|
if not is_qa_task and metric.startswith("qa_f1"):
|
|
continue
|
|
|
|
for length_group in all_length_groups:
|
|
if length_group == "overall":
|
|
# Look for metric without length suffix
|
|
metric_key = f"{split_key}_{metric}"
|
|
column_name = f"{task}_{metric}_overall"
|
|
else:
|
|
# Look for metric with length suffix
|
|
metric_key = f"{split_key}_{metric}_len_{length_group}"
|
|
column_name = f"{task}_{metric}_{length_group}"
|
|
|
|
value = metrics.get(metric_key, "N/A")
|
|
# Handle None values and convert to string
|
|
if value is None or value == "None":
|
|
value = "N/A"
|
|
elif isinstance(value, (int, float)):
|
|
value = f"{value:.4f}"
|
|
|
|
row_data[column_name] = value
|
|
|
|
# Create DataFrame with single row
|
|
if len(row_data) > 2: # More than just model and model_type
|
|
df = pd.DataFrame([row_data])
|
|
|
|
csv_filename = "evaluation_results"
|
|
if remove_context:
|
|
csv_filename += "_no_context"
|
|
csv_filename += ".csv"
|
|
csv_path = os.path.join(output_dir, csv_filename)
|
|
df.to_csv(csv_path, index=False)
|
|
print(f"Evaluation results saved to: {csv_path}")
|
|
else:
|
|
print("No evaluation data found to save to CSV")
|
|
|
|
|
|
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]
|
|
|
|
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
|
|
else:
|
|
rouge_metrics = compute_rouge(pred_texts, label_texts)
|
|
for k, v in rouge_metrics.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:
|
|
if ds_name in CLOSED_QA_DATASETS:
|
|
eval_result.metrics[f"{split_name}_qa_f1_len_{low}-{high}"] = "None"
|
|
else:
|
|
eval_result.metrics[f"{split_name}_rougeL.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:
|
|
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
|
|
else:
|
|
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
|
|
|
|
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()
|
|
|
|
# Create CSV summary of all evaluation results
|
|
if out:
|
|
# Determine model name and type based on the first model we can infer
|
|
model_name = "unknown_model"
|
|
is_hypernet = False
|
|
|
|
# Try to extract model name from the first split name or other context
|
|
if hasattr(eval_trainer.model, "base_model"):
|
|
if hasattr(eval_trainer.model.base_model, "config"):
|
|
model_name = getattr(
|
|
eval_trainer.model.base_model.config,
|
|
"name_or_path",
|
|
getattr(
|
|
eval_trainer.model.base_model.config, "_name_or_path", "unknown"
|
|
),
|
|
)
|
|
is_hypernet = hasattr(eval_trainer.model, "ctx_encoder")
|
|
elif hasattr(eval_trainer.model, "config"):
|
|
model_name = getattr(
|
|
eval_trainer.model.config,
|
|
"name_or_path",
|
|
getattr(eval_trainer.model.config, "_name_or_path", "unknown"),
|
|
)
|
|
|
|
# Clean up model name for display
|
|
if "/" in model_name:
|
|
model_name = model_name.split("/")[-1]
|
|
|
|
create_metrics_csv(
|
|
out,
|
|
output_dir=eval_trainer.args.output_dir,
|
|
model_name=model_name,
|
|
is_hypernet_model=is_hypernet,
|
|
remove_context=remove_context,
|
|
)
|
|
|
|
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,
|
|
)
|