|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| from typing import TYPE_CHECKING, Optional
|
|
|
| from ...data import KTODataCollatorWithPadding, get_dataset, get_template_and_fix_tokenizer
|
| from ...extras.constants import IGNORE_INDEX
|
| from ...extras.ploting import plot_loss
|
| from ...hparams import ModelArguments
|
| from ...model import load_model, load_tokenizer
|
| from ..trainer_utils import create_modelcard_and_push, create_ref_model
|
| from .trainer import CustomKTOTrainer
|
|
|
|
|
| if TYPE_CHECKING:
|
| from transformers import Seq2SeqTrainingArguments, TrainerCallback
|
|
|
| from ...hparams import DataArguments, FinetuningArguments
|
|
|
|
|
| def run_kto(
|
| model_args: "ModelArguments",
|
| data_args: "DataArguments",
|
| training_args: "Seq2SeqTrainingArguments",
|
| finetuning_args: "FinetuningArguments",
|
| callbacks: Optional[list["TrainerCallback"]] = None,
|
| ):
|
| tokenizer_module = load_tokenizer(model_args)
|
| tokenizer = tokenizer_module["tokenizer"]
|
| template = get_template_and_fix_tokenizer(tokenizer, data_args)
|
| dataset_module = get_dataset(template, model_args, data_args, training_args, stage="kto", **tokenizer_module)
|
| model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train)
|
|
|
| data_collator = KTODataCollatorWithPadding(
|
| template=template,
|
| model=model,
|
| pad_to_multiple_of=8,
|
| label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id,
|
| **tokenizer_module,
|
| )
|
|
|
|
|
| if finetuning_args.ref_model is None and (not training_args.do_train):
|
| ref_model = model
|
| else:
|
| ref_model = create_ref_model(model_args, finetuning_args)
|
|
|
|
|
| trainer = CustomKTOTrainer(
|
| model=model,
|
| ref_model=ref_model,
|
| args=training_args,
|
| finetuning_args=finetuning_args,
|
| data_collator=data_collator,
|
| callbacks=callbacks,
|
| **dataset_module,
|
| **tokenizer_module,
|
| )
|
|
|
|
|
| if training_args.do_train:
|
| train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
|
| trainer.save_model()
|
| trainer.log_metrics("train", train_result.metrics)
|
| trainer.save_metrics("train", train_result.metrics)
|
| trainer.save_state()
|
| if trainer.is_world_process_zero() and finetuning_args.plot_loss:
|
| keys = ["loss", "rewards/chosen"]
|
| if isinstance(dataset_module.get("eval_dataset"), dict):
|
| keys += [f"eval_{key}_loss" for key in dataset_module["eval_dataset"].keys()]
|
| else:
|
| keys += ["eval_loss"]
|
|
|
| plot_loss(training_args.output_dir, keys=keys)
|
|
|
|
|
| if training_args.do_eval:
|
| metrics = trainer.evaluate(metric_key_prefix="eval")
|
| if id(model) == id(ref_model):
|
| remove_keys = [key for key in metrics.keys() if "rewards" in key]
|
| for key in remove_keys:
|
| metrics.pop(key)
|
| trainer.log_metrics("eval", metrics)
|
| trainer.save_metrics("eval", metrics)
|
|
|
|
|
| create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)
|
|
|