fixing prefix_allowed_tokens_fn

[nicola-decao]

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
• Did you read the contributor guideline?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

Fixes the use of prefix_allowed_tokens_fn in generation. It was working for fairseq==0.9.0 (see https://github.com/facebookresearch/GENRE) but with the current version is broken.

PR review

Anyone in the community is free to review the PR once the tests have passed.

Did you have fun?

Make sure you had fun coding 🙃

[facebook-github-bot]

@myleott has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[erip]

I'm interested in seeing this land, but I was curious if it would be possible to also include a bit of documentation about this prefix_allowed_tokens_fn callable. I can't seem to find anything which explains what it's supposed to do or what shape it's supposed to be (f: (Tensor, list[int]) -> int?)

[nicola-decao]

I'm interested in seeing this land, but I was curious if it would be possible to also include a bit of documentation about this prefix_allowed_tokens_fn callable. I can't seem to find anything which explains what it's supposed to do or what shape it's supposed to be (f: (Tensor, list[int]) -> int?)

@erip you are right. Where do you think is the best place to write the signature of the function? here? https://github.com/pytorch/fairseq/blob/e5e8b3fee1e57a7abf35ad1a3ff223a2b7190c65/fairseq/search.py#L148

[erip]

@nicola-decao I think that makes good sense. There doesn't seem to be any documentation on the other search strategies, but this one is somewhat less straightforward since it's got the callback. Unless @myleott has other thoughts, I think throwing a docstring beneath the ctor would be great.

[myleott]

@nicola-decao if you can share a docstring here, I can update the imported version before merging

[nicola-decao]

@nicola-decao if you can share a docstring here, I can update the imported version before merging

@myleott Here you go:

prefix_allowed_tokens_fn: Callable[[int, torch.Tensor], List[int]]: If provided, this function constrains the beam search to allowed tokens only at each step. If not provided no constraint is applied. This function takes 2 arguments: the batch ID batch_id: int and a unidimensional tensor of token ids inputs_ids: torch.Tensor. It has to return a List[int] with the allowed tokens for the next generation step conditioned on the previously generated tokens inputs_ids and the batch ID batch_id. This argument is useful for constrained generation conditioned on the prefix, as described in Autoregressive Entity Retrieval https://arxiv.org/abs/2010.00904 and https://github.com/facebookresearch/GENRE.

[nicola-decao]

@myleott Any news on this? Is there something I should do?

[facebook-github-bot]

@myleott has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[nicola-decao]

@myleott Any news on this? I have a facebook AI project https://github.com/facebookresearch/GENRE that depends on this bug fix (for now I link people to my fork with the fix that is not ideal).

[nicola-decao]

@myleott @sshleifer can we please proceed on the merge here? It is really a minor change.

There is this Facebook AI project https://github.com/facebookresearch/GENRE that depends on this bug fix (for now I link people to my fork with the fix that is not ideal and may have trouble installing it).

[erip]

also cc @alexeib