Add support for prefix_allowed_tokens_fn to maintain a state throughout decoding
Author: John-Boccio
Pull request: no
Labels: name: Feature request

name: Generation

### Feature request

Add an optional argument in `prefix_allow_tokens_fn` to allow for state to maintained throughout decoding or add a stateful alternative to `prefix_allowed_tokens_fn`.

### Motivation

`prefix_allowed_tokens_fn` is great but has one major downfall which is that you cannot maintain a state throughout decoding. This is inefficient because at each step you must go through your past `inputIds`, build up your current "state", and then figure out which tokens are allowed to appear next.

Instead, there should be a class we can subclass that gets passed the next token ID at each step of decoding (`Constraint` does not achieve this as `update` does not get every token ID). For example if you are trying to create a function to output json format (https://gist.github.com/BorisTheBrave/969f303a082c9da1916d04ee1eb04452), then you could track where you currently on in the json as each token ID is being received instead of going through everything on each new token.

### Your contribution

Unfortunately can't make a PR.

Comments:
cc @gante 

Hi @John-Boccio 👋 

We don't need `transformers` changes to enable your request :D You can parameterize an arbitrary function with a mutable input and use that mutable input as the state in your `prefix_allowed_tokens_fn`.

Here's an example of how to create a function with state:
```py
from functools import partial

# A function with a `state` input, assumed to be a dictionary
def some_fn(foo, state=None):
    if state is None or not isinstance(state, dict):
        raise ValueError('`state` must be provided as a dictionary.')
    else:
        if 'bar' in state:
            state['bar'] += 1
        else:
            state['bar'] = 1
    return foo + state['bar']

# partial() allows us to create a new function from `some_fn` with a fixed value for `state`.
# Because the `state` input is mutable, the new function will keep track of the changes to `state`.
parameterized_fn = partial(some_fn, state={'bar': 0})

print(parameterized_fn(0))  # 1
print(parameterized_fn(0))  # 2
print(parameterized_fn(0))  # 3
```

Hi @gante ! Thank you for the suggestion. I actually had a similar idea and it does work well but has one catch - you must be performing greedy decoding. As soon as you add more than 1 beam, then all the beams will be sharing the objects passed into the partial function (i.e. all beams share the same state with no way to distinguish which beam you're operating on currently).

I think there will have to be some sort of new parameter to `generate` along the lines of `prefix_allowed_tokens_cls` which allows you to pass in a class that should be created for each beam that is used during generation.

For beam search to track specific beams, you would have to change a few things indeed -- including the API of the `LogitsProcessors`, to pass the previous beam indices so it could be passed to `prefix_allowed_tokens_fn`.

This falls outside the scope of what we want to support in `transformers`, at least for now 🤗 My suggestion would be to fork the library and change the generation loop to your needs :)