fix: preserve output logprobs without input logprobs (#34627)
Signed-off-by: jain-ria <riajain@NVIDIA.com>
This commit is contained in:
@@ -2655,11 +2655,9 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
||||
recv_obj: BatchStrOutput,
|
||||
recv_obj_index: int,
|
||||
):
|
||||
if recv_obj.input_token_logprobs_val is None:
|
||||
return
|
||||
|
||||
if (
|
||||
len(recv_obj.input_token_logprobs_val) > 0
|
||||
recv_obj.input_token_logprobs_val is not None
|
||||
and len(recv_obj.input_token_logprobs_val) > 0
|
||||
and recv_obj.input_token_logprobs_val[recv_obj_index] is not None
|
||||
):
|
||||
state.input_token_logprobs_val.extend(
|
||||
@@ -2668,12 +2666,16 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
||||
state.input_token_logprobs_idx.extend(
|
||||
recv_obj.input_token_logprobs_idx[recv_obj_index]
|
||||
)
|
||||
state.output_token_logprobs_val.extend(
|
||||
recv_obj.output_token_logprobs_val[recv_obj_index]
|
||||
)
|
||||
state.output_token_logprobs_idx.extend(
|
||||
recv_obj.output_token_logprobs_idx[recv_obj_index]
|
||||
)
|
||||
if (
|
||||
recv_obj.output_token_logprobs_val is not None
|
||||
and recv_obj.output_token_logprobs_val[recv_obj_index] is not None
|
||||
):
|
||||
state.output_token_logprobs_val.extend(
|
||||
recv_obj.output_token_logprobs_val[recv_obj_index]
|
||||
)
|
||||
state.output_token_logprobs_idx.extend(
|
||||
recv_obj.output_token_logprobs_idx[recv_obj_index]
|
||||
)
|
||||
|
||||
if top_logprobs_num > 0:
|
||||
if len(recv_obj.input_top_logprobs_val) > 0:
|
||||
|
||||
Reference in New Issue
Block a user