mnoukhov/pythia410m-sft-tldr
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: mnoukhov/pythia410m-sft-tldr
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-05-30 14:21:21 +08:00
commit 4bb3a5553f
215 changed files with 118927 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

471
code/callbacks.py Normal file
View File

@@ -0,0 +1,471 @@
import math
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union
import accelerate
import torch
from datasets import Dataset
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import PreTrainedTokenizerBase, TrainerCallback
import wandb
from trl.trainer.utils import pad_to_length
@dataclass
class PromptAndTextCollator:
tokenizer: PreTrainedTokenizerBase
padding: Union[bool, str] = True
max_prompt_length: Optional[int] = None
max_length: Optional[int] = None
prompt_field: str = "prompt"
target_field: str = "label"
return_tensors: str = "pt"
def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
prompts = [feat[self.prompt_field] for feat in features]
texts = [feat[self.prompt_field] + " " + feat[self.target_field] for feat in features]
original_side = self.tokenizer.padding_side
self.tokenizer.padding_side = "left"
tokenized_batch = self.tokenizer(
prompts,
truncation=True,
padding=True,
max_length=self.max_prompt_length,
return_tensors=self.return_tensors,
)
tokenized_batch["prompt"] = prompts
self.tokenizer.padding_side = original_side
tokenized_texts = self.tokenizer(
texts,
truncation=True,
padding=True,
max_length=self.max_length,
return_tensors=self.return_tensors,
)
text_labels = tokenized_texts["input_ids"].clone()
if self.tokenizer.pad_token_id is not None:
text_labels[text_labels == self.tokenizer.pad_token_id] = -100
tokenized_batch.update(
{
"text_input_ids": tokenized_texts["input_ids"],
"text_attention_mask": tokenized_texts["attention_mask"],
"text_labels": text_labels,
}
)
return tokenized_batch
class GoldModelRewardCallback(TrainerCallback):
def __init__(
self,
args,
gold_model,
gold_eval_dataset,
tokenizer,
accelerator,
max_length,
max_prompt_length,
prompt_field,
target_field,
gold_load_and_unload=False,
log_n_samples_during_eval=0,
generation_config=None,
):
self.max_length = max_length
self.log_n_samples_during_eval = log_n_samples_during_eval
self.generation_config = generation_config
# data_collator = DataCollatorWithPadding(tokenizer)
data_collator = PromptAndTextCollator(
tokenizer,
max_prompt_length=max_prompt_length,
max_length=max_length,
prompt_field=prompt_field,
target_field=target_field,
)
dataloader_params = {
"batch_size": args.eval_batch_size,
"collate_fn": data_collator,
"num_workers": args.dataloader_num_workers,
"pin_memory": args.dataloader_pin_memory,
}
dataloader = DataLoader(gold_eval_dataset, **dataloader_params)
self.dataloader = accelerator.prepare(dataloader)
self.accelerator = accelerator
self.completed_step = -1
self.gold_model = gold_model
self.gold_load_and_unload = gold_load_and_unload
# keep model on gpu the whole time
if not self.gold_load_and_unload:
self.gold_model = self.accelerator.prepare(self.gold_model)
def on_evaluate(self, args, state, control, model, tokenizer, metrics, **kwargs):
samples_to_log = []
gold_reward_sum = 0.0
nll_sum = 0.0
total_samples = 0
sample_length_sum = 0.0
# load model onto gpu for inference then unload
if self.gold_load_and_unload:
self.gold_model = self.accelerator.prepare(self.gold_model)
if state.global_step == self.completed_step:
return
for inputs in tqdm(
self.dataloader, desc="Gold Eval", dynamic_ncols=True, disable=not state.is_local_process_zero
):
# get loss over true continuation i.e. ppl on dataset
with torch.no_grad():
nll_loss = model(
input_ids=inputs["text_input_ids"],
attention_mask=inputs["text_attention_mask"],
labels=inputs["text_labels"],
).loss
nll_loss = self.accelerator.gather_for_metrics(nll_loss)
# generate from model
policy_output_decoded, ref_output_decoded, policy_output_ids = self.get_batch_samples(
model,
tokenizer,
inputs["input_ids"],
inputs["attention_mask"],
return_ids=True,
)
# gold reward
policy_output_attention_mask = (policy_output_ids != tokenizer.pad_token_id).to(torch.int64)
with torch.no_grad():
gold_rewards = self.gold_model(
input_ids=policy_output_ids, attention_mask=policy_output_attention_mask
)[0]
gold_rewards = self.accelerator.gather_for_metrics(gold_rewards)
if state.is_local_process_zero:
nll_sum += nll_loss.sum().item()
gold_reward_sum += gold_rewards.sum().item()
total_samples += gold_rewards.size(0)
sample_length_sum += policy_output_attention_mask.sum().item()
# Sample and save to game log if requested (for one batch to save time)
for i, (prompt, pol, ref) in enumerate(
zip(inputs["prompt"], policy_output_decoded, ref_output_decoded)
):
if len(samples_to_log) < self.log_n_samples_during_eval:
samples_to_log.append([prompt, pol[len(prompt) :], ref[len(prompt) :]])
else:
break
if self.gold_load_and_unload:
self.gold_model = self.gold_model.to("cpu")
torch.cuda.empty_cache()
if state.is_world_process_zero:
gold_log = {
"eval/gold_rewards_mean": gold_reward_sum / total_samples,
"eval/perplexity": math.exp(nll_sum / total_samples),
"eval/gold_sample_length": sample_length_sum / total_samples,
}
for key, value in gold_log.items():
print(f"{key}: {value}")
if state.epoch:
gold_log["epoch"] = round(state.epoch, 2)
gold_log["step"] = state.global_step
if samples_to_log:
gold_log["gold_log"] = (
wandb.Table(
columns=["Prompt", "Policy", "Ref Model"],
rows=samples_to_log,
),
)
wandb.log(gold_log)
self.completed_step = state.global_step
def get_batch_samples(self, model, tokenizer, input_ids, attention_mask, return_ids=False) -> Tuple[str, str]:
"""Reduce inputs to unseen prompts, and maximum batch size if necessary
Generate samples from the model and reference model for the given batch of inputs."""
policy_output = model.generate(
input_ids=input_ids,
attention_mask=attention_mask,
generation_config=self.generation_config,
)
# if self.ref_model is None:
with self.accelerator.unwrap_model(model).disable_adapter():
reference_output = model.generate(
input_ids=input_ids,
attention_mask=attention_mask,
generation_config=self.generation_config,
)
# else:
# reference_output = self.ref_model.generate(
# **inputs,
# generation_config=self.generation_config,
# )
policy_output = pad_to_length(policy_output, self.max_length, tokenizer.pad_token_id)
policy_output_decoded = tokenizer.batch_decode(policy_output, skip_special_tokens=True)
reference_output = pad_to_length(reference_output, self.max_length, tokenizer.pad_token_id)
reference_output_decoded = tokenizer.batch_decode(reference_output, skip_special_tokens=True)
if return_ids:
return policy_output_decoded, reference_output_decoded, policy_output
else:
return policy_output_decoded, reference_output_decoded
class PerplexityCallback(TrainerCallback):
"""Like GoldModelReward in that you generate and get ppl on dataset
But you don't run eval with the gold model
Useful when gold model is very larger and you want to run inference later
"""
def __init__(
self,
args,
dataset,
tokenizer,
accelerator,
max_length,
max_prompt_length,
prompt_field,
target_field,
hub_model_id=None,
**kwargs,
):
self.max_length = max_length
# data_collator = DataCollatorWithPadding(tokenizer)
data_collator = PromptAndTextCollator(
tokenizer,
max_prompt_length=max_prompt_length,
max_length=max_length,
prompt_field=prompt_field,
target_field=target_field,
)
dataloader_params = {
"batch_size": args.eval_batch_size,
"collate_fn": data_collator,
"num_workers": args.dataloader_num_workers,
"pin_memory": args.dataloader_pin_memory,
}
dataloader = DataLoader(dataset, **dataloader_params)
self.dataloader = accelerator.prepare(dataloader)
self.accelerator = accelerator
self.completed_step = -1
self.hub_model_id = hub_model_id
def on_evaluate(self, args, state, control, model, tokenizer, metrics, **kwargs):
nll_sum = 0.0
total_samples = 0
if state.global_step == self.completed_step:
return
for inputs in tqdm(
self.dataloader, desc="PPL and Gen Eval", dynamic_ncols=True, disable=not state.is_local_process_zero
):
# get loss over true continuation i.e. ppl on dataset
with torch.no_grad():
nll_loss = model(
input_ids=inputs["text_input_ids"],
attention_mask=inputs["text_attention_mask"],
labels=inputs["text_labels"],
).loss
nll_loss = self.accelerator.gather_for_metrics(nll_loss)
if state.is_local_process_zero:
total_samples += nll_loss.size(0)
nll_sum += nll_loss.sum().item()
if state.is_world_process_zero:
# gather_for_metrics doesn't work for list of strings?
gold_log = {
"eval/perplexity": math.exp(nll_sum / total_samples),
}
for key, value in gold_log.items():
print(f"{key}: {value}")
if state.epoch:
gold_log["epoch"] = round(state.epoch, 2)
gold_log["step"] = state.global_step
wandb.log(gold_log)
if self.hub_model_id is not None:
model.push_to_hub(self.hub_model_id, revision=f"step{state.global_step}")
self.completed_step = state.global_step
class PerplexityGenCallback(TrainerCallback):
"""Like GoldModelReward in that you generate and get ppl on dataset
But you don't run eval with the gold model
Useful when gold model is very larger and you want to run inference later
"""
def __init__(
self,
args,
dataset,
tokenizer,
accelerator,
max_length,
max_prompt_length,
prompt_field,
target_field,
log_n_samples_during_eval=0,
generation_config=None,
hub_model_id="tmp",
):
self.max_length = max_length
self.log_n_samples_during_eval = log_n_samples_during_eval
self.generation_config = generation_config
# data_collator = DataCollatorWithPadding(tokenizer)
data_collator = PromptAndTextCollator(
tokenizer,
max_prompt_length=max_prompt_length,
max_length=max_length,
prompt_field=prompt_field,
target_field=target_field,
)
dataloader_params = {
"batch_size": args.eval_batch_size,
"collate_fn": data_collator,
"num_workers": args.dataloader_num_workers,
"pin_memory": args.dataloader_pin_memory,
}
dataloader = DataLoader(dataset, **dataloader_params)
self.dataloader = accelerator.prepare(dataloader)
self.accelerator = accelerator
self.completed_step = -1
self.hub_name = hub_model_id
def on_evaluate(self, args, state, control, model, tokenizer, metrics, **kwargs):
all_generations = []
all_prompts = []
nll_sum = 0.0
total_samples = 0
sample_length_sum = 0.0
if state.global_step == self.completed_step:
return
for inputs in tqdm(
self.dataloader, desc="PPL and Gen Eval", dynamic_ncols=True, disable=not state.is_local_process_zero
):
# get loss over true continuation i.e. ppl on dataset
with torch.no_grad():
nll_loss = model(
input_ids=inputs["text_input_ids"],
attention_mask=inputs["text_attention_mask"],
labels=inputs["text_labels"],
).loss
# generate from model
policy_output_ids = model.generate(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
generation_config=self.generation_config,
)
policy_output_ids = pad_to_length(policy_output_ids, self.max_length, tokenizer.pad_token_id)
policy_output_attention_mask = (policy_output_ids != tokenizer.pad_token_id).to(torch.int64)
generation_sizes = policy_output_attention_mask.sum(dim=1)
(nll_loss, generation_ids, generation_sizes) = self.accelerator.gather_for_metrics(
(nll_loss, policy_output_ids, generation_sizes)
)
prompts = accelerate.utils.gather_object(inputs["prompt"])
if state.is_local_process_zero:
nll_sum += nll_loss.sum().item()
total_samples += generation_sizes.size(0)
sample_length_sum += generation_sizes.sum().item()
generation_strs = tokenizer.batch_decode(generation_ids, skip_special_tokens=True)
all_prompts.extend(prompts)
all_generations.extend(generation_strs)
if state.is_world_process_zero:
# gather_for_metrics doesn't work for list of strings?
gold_log = {
"eval/perplexity": math.exp(nll_sum / total_samples),
"eval/gold_sample_length": sample_length_sum / total_samples,
}
for key, value in gold_log.items():
print(f"{key}: {value}")
if state.epoch:
gold_log["epoch"] = round(state.epoch, 2)
gold_log["step"] = state.global_step
if self.log_n_samples_during_eval:
samples_to_log = [
[prompt, generation[len(prompt) :]]
for prompt, generation in zip(
all_prompts[: self.log_n_samples_during_eval],
all_generations[: self.log_n_samples_during_eval],
)
]
gold_log["gold_log"] = (
wandb.Table(
columns=["Prompt", "Policy"],
rows=samples_to_log,
),
)
wandb.log(gold_log)
generation_ds = Dataset.from_dict({"generations": all_generations})
generation_ds.push_to_hub(f"{self.hub_name}_generations", revision=str(state.global_step))
self.completed_step = state.global_step
def get_batch_samples(self, model, tokenizer, input_ids, attention_mask, return_ids=False) -> Tuple[str, str]:
"""Reduce inputs to unseen prompts, and maximum batch size if necessary
Generate samples from the model and reference model for the given batch of inputs."""
policy_output = model.generate(
input_ids=input_ids,
attention_mask=attention_mask,
generation_config=self.generation_config,
)
# if self.ref_model is None:
with self.accelerator.unwrap_model(model).disable_adapter():
reference_output = model.generate(
input_ids=input_ids,
attention_mask=attention_mask,
generation_config=self.generation_config,
)
# else:
# reference_output = self.ref_model.generate(
# **inputs,
# generation_config=self.generation_config,
# )
policy_output = pad_to_length(policy_output, self.max_length, tokenizer.pad_token_id)
policy_output_decoded = tokenizer.batch_decode(policy_output, skip_special_tokens=True)
reference_output = pad_to_length(reference_output, self.max_length, tokenizer.pad_token_id)
reference_output_decoded = tokenizer.batch_decode(reference_output, skip_special_tokens=True)
if return_ids:
return policy_output_decoded, reference_output_decoded, policy_output
else:
return policy_output_decoded, reference_output_decoded