Support token suppression, forced tokens (besides eos and bos), and decoder prompting for flax generation

[andyehrenberg]

Feature request

Add logits processors for token suppression and forced tokens at specific indices.
Enable prompting the decoder of encoder-decoder models with decoder_input_ids.

Motivation

Currently, the flax generation utilities do not support token suppression, forcing specific tokens to be decoded at specific response indices, nor prompting the decoder (helpful for models like Whisper that support decoder prompts - Flax Whisper is implemented in #20479). Adding these would move the flax utilities closer to feature parity with the pytorch generation utilities. Adding these features would fully unlock a flax implementation of Whisper inference.

Your contribution

I already have these features implemented in a branch of my fork - happy to open a PR!

[andyehrenberg]

@sanchit-gandhi @patil-suraj @gante

[sanchit-gandhi]

Did we decide to implement these features in the Flax Whisper PR in the end? cc @ArthurZucker

[andyehrenberg]

@sanchit-gandhi @ArthurZucker I just added these back into the Flax Whisper PR

[sanchit-gandhi]

Cool! Closing this issue in favour of the PR #20479