Open John-Boccio opened 7 months ago
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:
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 :)
Feature request
Add an optional argument in
prefix_allow_tokens_fn
to allow for state to maintained throughout decoding or add a stateful alternative toprefix_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 pastinputIds
, 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 asupdate
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.