Add support for returning alternative tokens #297
Conversation
JTS22
changed the title
JTS22
force-pushed
the
main
branch
2 times, most recently
from
446455f
to
2f714e6
Compare
JTS22
changed the title
| ) | ||
| alternative_token_texts.append(alternative_token_text) | ||
| all_input_ids.pop() | ||
| alternative_tokens = AlternativeTokens( |
There was a problem hiding this comment.
It looks like we're overriding this at every iteration of the loop. Should alternative_tokens instead be a list of AlternativeTokens, one for each request in the batch?
There was a problem hiding this comment.
We are adding the alternative_tokens to the Generation that is created for each request, so I don't see a problem with overriding it. Basically I just copied from what is done with the PrefillTokens. The variables that go into AlternativeTokens are kind of badly named though, so I added another commit to fix that :)
There was a problem hiding this comment.
Ah my mistake I overlooked where the Generation creation was happening in this loop when looking at the PR earlier. Looks good!
| ) | ||
| alternative_token_texts.append(alternative_token_text) | ||
| all_input_ids.pop() | ||
| alternative_tokens = AlternativeTokens( |
There was a problem hiding this comment.
Ah my mistake I overlooked where the Generation creation was happening in this loop when looking at the PR earlier. Looks good!
| request_alternative_token_logprobs = alternative_token_logprobs[i][:num_alternatives] | ||
|
|
||
| # Decode tokens | ||
| request_alternative_token_texts = list() |
There was a problem hiding this comment.
Super minor nit, but generally I prefer using [] to list() to avoid symbol lookup. Definitely not a blocker.
Add support for returning alternative tokens
Fixes #298
So far, alternative tokens are only supported on
FlashCausalLMbased models.Before submitting
to it if that's the case: Include logprobs for alternative, less probable tokens in the generation response #298
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.