T5 generate only words of the input

Beginners
1

Hi, I want T5 to only be allowed to generate words of the current input sentence in the batch. I tried assigning batch_input_ids to prefix_allowed_tokens_fn but it doesn’t work properly.
Thanks in advance.