Example of prefix_allowed_tokens_fn() while text generation

Hello, I would like to use the prefix_allowed_tokens_fn as an input to the model.generate() function, in order to perform constrained text generation with BART.

I tried to adapt the function in the original repository here, but it doesn’t seem to be working. Can you please tell me if there are any examples of the kinds of functions that can be given as input to this parameter? Thank you!