Former-commit-id: d86455f685fa531e651333e00b4fe54d895cf2e4
This commit is contained in:
hiyouga
2024-01-09 18:31:27 +08:00
parent 89f4ae51f9
commit 4d6669c268
9 changed files with 78 additions and 50 deletions

View File

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Optional, List
from transformers import Seq2SeqTrainingArguments
from llmtuner.data import get_dataset, preprocess_dataset, split_dataset
from llmtuner.extras.callbacks import SavePeftModelCallback
from llmtuner.extras.callbacks import FixValueHeadModelCallback
from llmtuner.extras.ploting import plot_loss
from llmtuner.model import load_model_and_tokenizer
from llmtuner.train.rm.collator import PairwiseDataCollatorWithPadding
@@ -40,7 +40,7 @@ def run_rm(
args=training_args,
tokenizer=tokenizer,
data_collator=data_collator,
callbacks=callbacks + [SavePeftModelCallback()],
callbacks=callbacks + [FixValueHeadModelCallback()],
compute_metrics=compute_accuracy,
**split_dataset(dataset, data_args, training_args)
)