Move chat template formatting inside trainer

huggingface · Sep 23, 2024 · 13b5ed0 · 13b5ed0
1 parent fc9a2f9
commit 13b5ed0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 10 deletions.
diff --git a/examples/scripts/reward_modeling.py b/examples/scripts/reward_modeling.py
@@ -51,7 +51,6 @@
 import warnings
 
 import torch
-from accelerate import PartialState
 from datasets import load_dataset
 from tqdm import tqdm
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, HfArgumentParser
@@ -63,7 +62,6 @@
     get_kbit_device_map,
     get_peft_config,
     get_quantization_config,
-    maybe_apply_chat_template,
     setup_chat_format,
 )
 from trl.commands.cli_utils import RewardScriptArguments
@@ -110,16 +108,11 @@
             " Make sure to pass --lora_task_type SEQ_CLS when using this script with PEFT."
         )
 
-    #############################
-    # Load and preprocess dataset
-    #############################
+    ##############
+    # Load dataset
+    ##############
     dataset = load_dataset(args.dataset_name)
 
-    with PartialState().local_main_process_first():
-        dataset = dataset.map(
-            maybe_apply_chat_template, num_proc=training_args.dataset_num_proc, fn_kwargs={"tokenizer": tokenizer}
-        )
-
     ##########
     # Training
     ##########

diff --git a/trl/trainer/reward_trainer.py b/trl/trainer/reward_trainer.py
@@ -30,6 +30,7 @@
 from transformers.trainer_utils import EvalPrediction
 from transformers.utils import is_peft_available
 
+from ..data_utils import maybe_apply_chat_template
 from .reward_config import RewardConfig
 from .utils import (
     RewardDataCollatorWithPadding,
@@ -229,6 +230,7 @@ def __init__(
         if "input_ids" not in train_dataset.column_names:
             with PartialState().local_main_process_first():
                 fn_kwargs = {"tokenizer": tokenizer}
+                train_dataset = train_dataset.map(maybe_apply_chat_template, fn_kwargs={"tokenizer": tokenizer})
                 train_dataset = train_dataset.map(
                     _tokenize,
                     batched=True,
@@ -243,6 +245,7 @@ def __init__(
                     num_proc=args.dataset_num_proc,
                 )
                 if eval_dataset is not None:
+                    eval_dataset = eval_dataset.map(maybe_apply_chat_template, fn_kwargs={"tokenizer": tokenizer})
                     eval_dataset = eval_dataset.map(
                         _tokenize,
                         fn_kwargs=fn_kwargs,