-
-
Notifications
You must be signed in to change notification settings - Fork 5.3k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[Core] Adding token ranks along with logprobs #3516
Conversation
@@ -447,6 +447,9 @@ def _sample( | |||
] | |||
return sample_results | |||
|
|||
def _get_ranks(x: torch.Tensor, indices: List[int]) -> torch.Tensor: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can you add docstring to this func? Like the shape of x tensors, and what exactly indices mean?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
will do!
@@ -17,9 +17,9 @@ | |||
class Logprob: | |||
"""Infos for supporting OpenAI compatible logprobs.""" | |||
logprob: float | |||
rank: Optional[int] = None |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can you comment what this means here
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
will do!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks @SwapnilDreams100!
batched_ranks_query_result = _get_ranks( | ||
logprobs[batched_logprobs_query_seq_indices], | ||
batched_logprobs_query_token_indices) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
let's move this after line 526, since we'll have to force a CPU<->GPU sync here anyway.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
will do!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks good, please move the invocation as outlined in the comment before merge!
One question I guess is whether these should also be exposed in the openai API responses (but doesn't necessarily need to be addressed in this PR). |
Co-authored-by: Swapnil Parekh <[email protected]>
Co-authored-by: Swapnil Parekh <[email protected]>
Adds the token rank functionality (for both prompt tokens and sample tokens) within the
Logprobs
object, by adding therank
property toLogprobs
.This PR also adds a couple of tests to verify the ranks in greedy and non-greedy settings.
The implementation here is meant to align with functionality in https://github.com/IBM/text-generation-inference (IBM's fork of HF TGI).