Pad prompts to the right in T5 examples and add EOS token to seq2seq prompts (#422)

* Pad prompts to the right in T5 examples. Add EOS token to prompts for seq2seq models like T5

* style: satisfy black

* fix(offline_pipeline): default `add_special_tokens` to `False`

* style: satisfy flake

---------

Co-authored-by: reciprocated <[email protected]>

Author: mikljohansson

examples/ilql_sentiments_t5.py

@@ -41,6 +41,7 @@ def get_positive_score(scores):
     ),
     tokenizer=TokenizerConfig(
         tokenizer_path="lvwerra/t5-imdb",
+        padding_side="right",
         truncation_side="right",
     ),
     optimizer=OptimizerConfig(

examples/ppo_sentiments_t5.py

@@ -43,6 +43,7 @@ def get_positive_score(scores):
     ),
     tokenizer=TokenizerConfig(
         tokenizer_path="lvwerra/t5-imdb",
+        padding_side="right",
         truncation_side="right",
     ),
     optimizer=OptimizerConfig(

examples/ppo_translation_t5.py

@@ -54,6 +54,7 @@
     ),
     tokenizer=TokenizerConfig(
         tokenizer_path="t5-large",
+        padding_side="right",
         truncation_side="right",
     ),
     optimizer=OptimizerConfig(

trlx/pipeline/offline_pipeline.py

@@ -117,10 +117,16 @@ class PromptPipeline(BasePipeline):
         max_prompt_length (`int`): max length of the prompt, if exceeded the prompt will be truncated according to
             tokenizer's truncation setting.
         tokenizer (`transformers.PreTrainedTokenizer`): a tokenizer to tokenize prompts with.
+        add_special_tokens (`bool`): whether to encode prompts with tokenizer's special tokens (passed directly
+            into `tokenizer.encode`)
     """
 
     def __init__(
-        self, prompts: Union[Dict[str, Any], List[str]], max_prompt_length: int, tokenizer: PreTrainedTokenizer
+        self,
+        prompts: Union[Dict[str, Any], List[str]],
+        max_prompt_length: int,
+        tokenizer: PreTrainedTokenizer,
+        add_special_tokens: bool = False,
     ):
         super().__init__()
 
@@ -131,7 +137,7 @@ def __init__(
             metadata = [{}] * len(prompts)
 
         model_inputs = tokenizer(
-            prompts, truncation=True, padding=False, max_length=max_prompt_length, add_special_tokens=False
+            prompts, truncation=True, padding=False, max_length=max_prompt_length, add_special_tokens=add_special_tokens
         )
 
         prompts_tokens = model_inputs["input_ids"]

trlx/pipeline/ppo_pipeline.py

@@ -15,10 +15,11 @@ class PPORolloutStorage(BaseRolloutStore):
     Rollout storage for training PPO
     """
 
-    def __init__(self, pad_token_id):
+    def __init__(self, pad_token_id, padding_side):
         super().__init__()
 
         self.pad_token_id = pad_token_id
+        self.padding_side = padding_side
         self.history: Iterable[PPORLElement] = [None]
 
     def push(self, exps: Iterable[PPORLElement]):
@@ -51,13 +52,23 @@ def create_loader(
         shuffle: bool,
     ) -> DataLoader:
         def collate_fn(elems: Iterable[PPORLElement]):
-            return PPORLBatch(
+            if self.padding_side == "right":
+                # Right padding of already right-padded queries
+                query_tensors = pad_sequence(
+                    [elem.query_tensor for elem in elems],
+                    padding_value=self.pad_token_id,
+                    batch_first=True,
+                )
+            else:
                 # Left padding of already left-padded queries
-                pad_sequence(
+                query_tensors = pad_sequence(
                     [elem.query_tensor.flip(0) for elem in elems],
                     padding_value=self.pad_token_id,
                     batch_first=True,
-                ).flip(1),
+                ).flip(1)
+
+            return PPORLBatch(
+                query_tensors,
                 # Right pad the rest, to have a single horizontal query/response split
                 pad_sequence(
                     [elem.response_tensor for elem in elems],

trlx/trainer/accelerate_ppo_trainer.py

@@ -54,7 +54,7 @@ def __init__(self, config: TRLConfig, **kwargs):
 
         # Setup the rollout store
         # Rollouts contain the prompt & response, log probs, values and rewards - from each rollout
-        self.store = PPORolloutStorage(self.tokenizer.pad_token_id)
+        self.store = PPORolloutStorage(self.tokenizer.pad_token_id, self.tokenizer.padding_side)
 
         # Create the rollout store dataloader (for batching up rollouts)
         # TODO (jon-tow): This is only used to satisfy to `accelerator.prepare` call constraint below - remove in future

trlx/trlx.py

@@ -94,7 +94,9 @@ def train(  # noqa: C901
         if eval_prompts is None:
             eval_prompts = prompts[:batch_size]
 
-        pipeline = get_pipeline(config.train.pipeline)(prompts, max_prompt_length, trainer.tokenizer)
+        pipeline = get_pipeline(config.train.pipeline)(
+            prompts, max_prompt_length, trainer.tokenizer, add_special_tokens=config.model.model_arch_type == "seq2seq"
+        )
         trainer.add_prompt_pipeline(pipeline)
 
         if eval_prompts is None:
@@ -118,7 +120,9 @@ def train(  # noqa: C901
     else:
         raise ValueError("Either `samples` or `reward_fn` should be given for training")
 
-    eval_pipeline = get_pipeline(config.train.pipeline)(eval_prompts, max_prompt_length, trainer.tokenizer)
+    eval_pipeline = get_pipeline(config.train.pipeline)(
+        eval_prompts, max_prompt_length, trainer.tokenizer, add_special_tokens=config.model.model_arch_type == "seq2seq"
+    )
     trainer.add_eval_pipeline(eval_pipeline)
 
     trainer.learn()