Fix PPO example (#4556)

Author: qgallouedec

examples/scripts/ppo/ppo_tldr.py

@@ -44,7 +44,7 @@
 
 """
 python examples/scripts/ppo/ppo_tldr.py \
-    --dataset_name trl-internal-testing/tldr-preference-sft-trl-style \
+    --dataset_name trl-lib/tldr \
     --dataset_test_split validation \
     --learning_rate 3e-6 \
     --output_dir pythia-1b-deduped-tldr-preference-sft-trl-style-ppo \
@@ -62,7 +62,7 @@
 
 accelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \
     examples/scripts/ppo/ppo_tldr.py \
-    --dataset_name trl-internal-testing/tldr-preference-sft-trl-style \
+    --dataset_name trl-lib/tldr \
     --dataset_test_split validation \
     --output_dir pythia-1b-deduped-tldr-preference-sft-trl-style-ppo \
     --learning_rate 3e-6 \
@@ -134,11 +134,7 @@ def prepare_dataset(dataset, tokenizer):
         """pre-tokenize the dataset before training; only collate during training"""
 
         def tokenize(element):
-            input_ids = tokenizer.apply_chat_template(
-                element["messages"][:1],
-                padding=False,
-                add_generation_prompt=True,
-            )
+            input_ids = tokenizer(element["prompt"], padding=False)["input_ids"]
             return {"input_ids": input_ids, "lengths": len(input_ids)}
 
         return dataset.map(