Fix rlhf on trl v0.11.0

satyaog · Sep 23, 2024 · 6caac29 · 6caac29
1 parent dbebc7c
commit 6caac29
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/benchmarks/rlhf/main.py b/benchmarks/rlhf/main.py
@@ -13,7 +13,7 @@
 
 from trl import ModelConfig
 from trl.trainer.ppov2_trainer import PPOv2Config, PPOv2Trainer
-from trl.trainer.utils import SIMPLE_QUERY_CHAT_TEMPLATE
+from trl.trainer.utils import SIMPLE_CHAT_TEMPLATE
 
 
 class PPOv2TrainerIntrumented(PPOv2Trainer):
@@ -62,7 +62,7 @@ def main():
     )
     tokenizer.add_special_tokens({"pad_token": "[PAD]"})
     if tokenizer.chat_template is None:
-        tokenizer.chat_template = SIMPLE_QUERY_CHAT_TEMPLATE
+        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE
     value_model = AutoModelForSequenceClassification.from_pretrained(
         config.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1
     )

diff --git a/benchmarks/rlhf/prepare.py b/benchmarks/rlhf/prepare.py
@@ -11,7 +11,7 @@
 from datasets import load_dataset
 from trl import ModelConfig
 from trl.trainer.ppov2_trainer import PPOv2Config
-from trl.trainer.utils import SIMPLE_QUERY_CHAT_TEMPLATE
+from trl.trainer.utils import SIMPLE_CHAT_TEMPLATE
 
 
 if __name__ == "__main__":
@@ -30,7 +30,7 @@
     tokenizer.add_special_tokens({"pad_token": "[PAD]"})
 
     if tokenizer.chat_template is None:
-        tokenizer.chat_template = SIMPLE_QUERY_CHAT_TEMPLATE
+        tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE
 
     value_model = AutoModelForSequenceClassification.from_pretrained(
         config.reward_model_path,