初始化项目，由ModelHub XC社区提供模型
Model: mnoukhov/pythia410m-sft-tldr Source: Original Platform
This commit is contained in:
ModelHub XC
216
code/evaluate_rouge.py
Normal file
216
code/evaluate_rouge.py
Normal file
@@ -0,0 +1,216 @@
|
||||
import shutil
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
import evaluate
|
||||
import nltk
|
||||
import torch
|
||||
from accelerate import Accelerator
|
||||
from datasets import load_dataset
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
from transformers import (
|
||||
AutoModelForCausalLM,
|
||||
AutoTokenizer,
|
||||
BitsAndBytesConfig,
|
||||
HfArgumentParser,
|
||||
default_data_collator,
|
||||
)
|
||||
|
||||
import wandb
|
||||
|
||||
|
||||
shutil.disk_usage = lambda x: shutil._ntuple_diskusage(1, 1, 1)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScriptArguments:
|
||||
model_name: Optional[str] = field(default="EleutherAI/pythia-6.9b-deduped", metadata={"help": "the model name"})
|
||||
tokenizer_name: Optional[str] = field(default=None, metadata={"help": "the tokenizer name"})
|
||||
dataset_name: Optional[str] = field(
|
||||
default="CarperAI/openai_summarize_tldr", metadata={"help": "the dataset name"}
|
||||
)
|
||||
split: Optional[str] = field(default="valid[:20]", metadata={"help": "the dataset name"})
|
||||
dataset_text_field: Optional[str] = field(default="prompt")
|
||||
dataset_label_field: Optional[str] = field(default="label")
|
||||
load_in_8bit: Optional[bool] = field(default=True, metadata={"help": "load the model in 8 bits precision"})
|
||||
load_in_4bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 4 bits precision"})
|
||||
use_peft: Optional[bool] = field(default=True, metadata={"help": "Wether to use PEFT or not to train adapters"})
|
||||
seed: Optional[int] = field(default=0)
|
||||
batch_size: Optional[int] = field(default=1)
|
||||
bf16: Optional[bool] = field(default=False)
|
||||
fp16: Optional[bool] = field(default=False)
|
||||
seq_length: Optional[int] = field(default=512, metadata={"help": "Input sequence length"})
|
||||
max_new_tokens: Optional[int] = field(default=48, metadata={"help": "Max new tokens to generate"})
|
||||
num_logged_samples: int = field(default=100, metadata={"help": "Max samples to log to wandb"})
|
||||
temperature: Optional[float] = field(default=0.0)
|
||||
sample: Optional[bool] = field(default=False)
|
||||
|
||||
strip: Optional[bool] = field(default=False)
|
||||
log_with: Optional[str] = field(default="wandb")
|
||||
|
||||
|
||||
parser = HfArgumentParser(ScriptArguments)
|
||||
args = parser.parse_args_into_dataclasses()[0]
|
||||
|
||||
print("Loading the model")
|
||||
if args.load_in_8bit and args.load_in_4bit:
|
||||
raise ValueError("You can't load the model in 8 bits and 4 bits at the same time")
|
||||
elif args.load_in_8bit or args.load_in_4bit:
|
||||
quantization_config = BitsAndBytesConfig(load_in_8bit=args.load_in_8bit, load_in_4bit=args.load_in_4bit)
|
||||
device_map = {"": Accelerator().local_process_index}
|
||||
else:
|
||||
device_map = None
|
||||
quantization_config = None
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
args.model_name,
|
||||
torch_dtype=torch.bfloat16 if args.bf16 else None,
|
||||
)
|
||||
|
||||
print("Loading dataset")
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
args.model_name if args.tokenizer_name is None else args.tokenizer_name, padding_side="left"
|
||||
)
|
||||
if getattr(tokenizer, "pad_token", None) is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
|
||||
def create_dataset(tokenizer, args):
|
||||
eval_data = load_dataset(
|
||||
args.dataset_name,
|
||||
split=args.split,
|
||||
)
|
||||
|
||||
padding = "max_length"
|
||||
max_source_length = args.seq_length
|
||||
max_target_length = 52
|
||||
|
||||
def preprocess_function(example):
|
||||
inputs = example[args.dataset_text_field]
|
||||
targets = example[args.dataset_label_field]
|
||||
|
||||
if args.strip:
|
||||
inputs = inputs.strip()
|
||||
targets = targets.strip()
|
||||
|
||||
model_inputs = tokenizer(inputs, max_length=max_source_length, padding=padding, truncation=True)
|
||||
|
||||
# Tokenize targets with the `text_target` keyword argument
|
||||
labels = tokenizer(text_target=targets, max_length=max_target_length, padding=padding, truncation=True)
|
||||
|
||||
model_inputs["labels"] = labels["input_ids"]
|
||||
return model_inputs
|
||||
|
||||
eval_dataset = eval_data.map(
|
||||
preprocess_function,
|
||||
remove_columns=eval_data.column_names,
|
||||
)
|
||||
|
||||
return eval_dataset
|
||||
|
||||
|
||||
eval_dataset = create_dataset(tokenizer, args)
|
||||
|
||||
rouge = evaluate.load("rouge")
|
||||
|
||||
|
||||
def postprocess_text(preds, labels):
|
||||
preds = [pred.strip() for pred in preds]
|
||||
labels = [label.strip() for label in labels]
|
||||
|
||||
# rougeLSum expects newline after each sentence
|
||||
preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds]
|
||||
labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels]
|
||||
|
||||
return preds, labels
|
||||
|
||||
|
||||
accelerator = Accelerator()
|
||||
eval_dataloader = DataLoader(eval_dataset, collate_fn=default_data_collator, batch_size=args.batch_size)
|
||||
|
||||
model, eval_dataloader = accelerator.prepare(model, eval_dataloader)
|
||||
|
||||
model.eval()
|
||||
|
||||
gen_kwargs = {
|
||||
"max_new_tokens": args.max_new_tokens,
|
||||
"pad_token_id": tokenizer.pad_token_id,
|
||||
"eos_token_id": tokenizer.eos_token_id,
|
||||
"do_sample": args.sample,
|
||||
}
|
||||
|
||||
wandb.init(project="trl")
|
||||
table = wandb.Table(columns=["prompt", "prediction", "label"])
|
||||
|
||||
for batch in tqdm(eval_dataloader):
|
||||
with torch.no_grad():
|
||||
output_tokens = accelerator.unwrap_model(model).generate(
|
||||
batch["input_ids"],
|
||||
attention_mask=batch["attention_mask"],
|
||||
**gen_kwargs,
|
||||
)
|
||||
|
||||
# get just the generated tokens
|
||||
generated_tokens = output_tokens[:, batch["input_ids"].shape[1] :]
|
||||
|
||||
generated_tokens = accelerator.pad_across_processes(generated_tokens, dim=1, pad_index=tokenizer.pad_token_id)
|
||||
labels = batch["labels"]
|
||||
# if not args.pad_to_max_length:
|
||||
# # If we did not pad to max length, we need to pad the labels too
|
||||
# labels = accelerator.pad_across_processes(batch["labels"], dim=1, pad_index=tokenizer.pad_token_id)
|
||||
|
||||
generated_tokens, labels = accelerator.gather_for_metrics((generated_tokens, labels))
|
||||
generated_tokens = generated_tokens.cpu().numpy()
|
||||
labels = labels.cpu().numpy()
|
||||
|
||||
# if args.ignore_pad_token_for_loss:
|
||||
# # Replace -100 in the labels as we can't decode them.
|
||||
# labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
|
||||
if isinstance(generated_tokens, tuple):
|
||||
generated_tokens = generated_tokens[0]
|
||||
decoded_preds = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
|
||||
decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
|
||||
|
||||
decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
|
||||
|
||||
# print(f"Label {decoded_labels}")
|
||||
# print(f"Pred {decoded_preds}")
|
||||
|
||||
rouge.add_batch(
|
||||
predictions=decoded_preds,
|
||||
references=decoded_labels,
|
||||
)
|
||||
|
||||
# log samples
|
||||
if accelerator.is_main_process:
|
||||
if len(table.data) < args.num_logged_samples:
|
||||
num_samples_to_add = min(args.num_logged_samples - len(table.data), len(batch["input_ids"]))
|
||||
for i in range(num_samples_to_add):
|
||||
table.add_data(
|
||||
tokenizer.decode(batch["input_ids"][i], skip_special_tokens=True),
|
||||
decoded_preds[i],
|
||||
decoded_labels[i],
|
||||
)
|
||||
|
||||
if len(table.data) == args.num_logged_samples:
|
||||
if args.log_with == "wandb":
|
||||
wandb.log({"examples": table})
|
||||
else:
|
||||
for row in table.iterrows():
|
||||
print("PROMPT")
|
||||
print(row[1][0])
|
||||
print("\n")
|
||||
print("PRED")
|
||||
print(row[1][1])
|
||||
print("\n")
|
||||
print("LABEL")
|
||||
print(row[1][2])
|
||||
print("\n")
|
||||
print("\n")
|
||||
|
||||
result = rouge.compute()
|
||||
print(result)
|
||||
|
||||
for key, value in result.items():
|
||||
wandb.run.summary[key] = value
|
||||
Reference in New Issue
Block a user