Repository: huggingface/trl

Issue:
why not use bacth for model generate
In [gpt2 example code](https://github.com/lvwerra/trl/blob/main/examples/sentiment/scripts/gpt2-sentiment.py#L177)，model generate text single by single，it make traing slow，why not use batch generate?


Discussion:
This is addressed in #228! We need to update the examples though!

Resolving patch:
diff --git a/examples/sentiment/scripts/gpt2-sentiment.py b/examples/sentiment/scripts/gpt2-sentiment.py
--- a/examples/sentiment/scripts/gpt2-sentiment.py
+++ b/examples/sentiment/scripts/gpt2-sentiment.py
@@ -173,13 +173,10 @@ def collator(data):
     query_tensors = batch["input_ids"]
 
     # Get response from gpt2
-    response_tensors = []
-    for query in query_tensors:
-        gen_len = output_length_sampler()
-        generation_kwargs["max_new_tokens"] = gen_len
-        response = ppo_trainer.generate(query, **generation_kwargs)
-        response_tensors.append(response.squeeze()[-gen_len:])
-    batch["response"] = [tokenizer.decode(r.squeeze()) for r in response_tensors]
+    response_tensors = ppo_trainer.generate(
+        query_tensors, return_prompt=False, length_sampler=output_length_sampler, **generation_kwargs
+    )
+    batch["response"] = tokenizer.batch_decode(response_tensors)
 
     # Compute sentiment score
     texts = [q + r for q, r in zip(batch["query"], batch["response"])]
diff --git a/examples/sentiment/scripts/gpt2-sentiment_peft.py b/examples/sentiment/scripts/gpt2-sentiment_peft.py
--- a/examples/sentiment/scripts/gpt2-sentiment_peft.py
+++ b/examples/sentiment/scripts/gpt2-sentiment_peft.py
@@ -214,13 +214,10 @@ def print_trainable_parameters(model):
     model.gradient_checkpointing_disable()
     model.pretrained_model.config.use_cache = True
     # Get response from Causal LM
-    response_tensors = []
-    for query in query_tensors:
-        gen_len = output_length_sampler()
-        generation_kwargs["max_new_tokens"] = gen_len
-        response = ppo_trainer.generate(query, **generation_kwargs)
-        response_tensors.append(response.squeeze()[-gen_len:])
-    batch["response"] = [tokenizer.decode(r.squeeze()) for r in response_tensors]
+    response_tensors = ppo_trainer.generate(
+        query_tensors, return_prompt=False, length_sampler=output_length_sampler, **generation_kwargs
+    )
+    batch["response"] = tokenizer.batch_decode(response_tensors)
 
     # Compute sentiment score
     texts = [q + r for q, r in zip(batch["query"], batch["response"])]
diff --git a/examples/sentiment/scripts/t5-sentiment.py b/examples/sentiment/scripts/t5-sentiment.py
--- a/examples/sentiment/scripts/t5-sentiment.py
+++ b/examples/sentiment/scripts/t5-sentiment.py
@@ -144,14 +144,12 @@ def collater(data):
 for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
     query_tensors = batch["input_ids"]
 
-    # Get response from gpt2
-    response_tensors = []
-    for query in query_tensors:
-        gen_len = output_length_sampler()
-        generation_kwargs["max_new_tokens"] = gen_len
-        response = ppo_trainer.generate(query, **generation_kwargs)
-        response_tensors.append(response.squeeze())
-    batch["response"] = [tokenizer.decode(r[1:].squeeze()) for r in response_tensors]
+    # Get response from t5
+    response_tensors = ppo_trainer.generate(
+        query_tensors, return_prompt=False, length_sampler=output_length_sampler, **generation_kwargs
+    )
+    response_tensors = [r[1:] for r in response_tensors]
+    batch["response"] = tokenizer.batch_decode(response_tensors)
 
     # Compute sentiment score
     texts = [q + r for q, r in zip(batch["query"], batch["response"])]


Test patch: