Using "beam search" strategy while generating the responses #2534
Open
Description
Hi
I am using flan-t5-xl to generate the output.
When I use the function ppo_trainer.generate(....), it gives me the desired output but I guess it is the top beam or the best output.
I am trying to generate output for 4 beams (currently using this custom generate function):
def select(list_data, model, tokenizer, strategy = "beam"):
device = 0 if torch.cuda.is_available() else "cpu"
beams, num_seq = 5, 4
batch = tokenizer(list_data, return_tensors="pt", padding=True)
if strategy == "beam":
generated = model.generate(
input_ids = batch["input_ids"].to(device), attention_mask = batch["attention_mask"].to(device),
num_beams=beams, early_stopping=True, num_return_sequences=num_seq, max_new_tokens=60, min_length=5
)
final_list = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True).lower() for g in generated]
return final_list
Is it possible to generate output for multiple beams using the ppo_trainer.generate(....) function?