Simplify stream_output (#2398)
This commit is contained in:
Lianmin Zheng
@@ -39,10 +39,12 @@ class LogitsProcessorOutput:
|
||||
# The logprobs of input tokens. shape: [#token, vocab_size]
|
||||
input_token_logprobs: torch.Tensor = None
|
||||
|
||||
# The logprob and id of the top-k tokens in input positions. shape [#seq, #token, k] of Tuple(logprob, token_id)
|
||||
input_top_logprobs: List = None
|
||||
# The logprob and id of the top-k tokens in output positions. shape [#seq, #token, k] of Tuple(logprob, token_id)
|
||||
output_top_logprobs: List = None
|
||||
# The logprob and id of the top-k tokens in input positions. shape [#seq, #token, k]
|
||||
input_top_logprobs_val: List = None
|
||||
input_top_logprobs_idx: List = None
|
||||
# The logprob and id of the top-k tokens in output positions. shape [#seq, #token, k]
|
||||
output_top_logprobs_val: List = None
|
||||
output_top_logprobs_idx: List = None
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
@@ -125,12 +127,15 @@ class LogitsProcessor(nn.Module):
|
||||
indices = ret.indices.tolist()
|
||||
|
||||
if logits_metadata.forward_mode.is_decode():
|
||||
output_top_logprobs = []
|
||||
output_top_logprobs_val = []
|
||||
output_top_logprobs_idx = []
|
||||
for i, k in enumerate(logits_metadata.top_logprobs_nums):
|
||||
output_top_logprobs.append(list(zip(values[i][:k], indices[i][:k])))
|
||||
return None, output_top_logprobs
|
||||
output_top_logprobs_val.append(values[i][:k])
|
||||
output_top_logprobs_idx.append(indices[i][:k])
|
||||
return None, None, output_top_logprobs_val, output_top_logprobs_idx
|
||||
else:
|
||||
input_top_logprobs, output_top_logprobs = [], []
|
||||
input_top_logprobs_val, input_top_logprobs_idx = [], []
|
||||
output_top_logprobs_val, output_top_logprobs_idx = [], []
|
||||
|
||||
pt = 0
|
||||
for k, pruned_len in zip(
|
||||
@@ -138,27 +143,36 @@ class LogitsProcessor(nn.Module):
|
||||
logits_metadata.extend_logprob_pruned_lens_cpu,
|
||||
):
|
||||
if pruned_len <= 0:
|
||||
input_top_logprobs.append([])
|
||||
output_top_logprobs.append([])
|
||||
input_top_logprobs_val.append([])
|
||||
input_top_logprobs_idx.append([])
|
||||
output_top_logprobs_val.append([])
|
||||
output_top_logprobs_idx.append([])
|
||||
continue
|
||||
|
||||
input_top_logprobs.append(
|
||||
[
|
||||
list(zip(values[pt + j][:k], indices[pt + j][:k]))
|
||||
for j in range(pruned_len - 1)
|
||||
]
|
||||
input_top_logprobs_val.append(
|
||||
[values[pt + j][:k] for j in range(pruned_len - 1)]
|
||||
)
|
||||
output_top_logprobs.append(
|
||||
input_top_logprobs_idx.append(
|
||||
[indices[pt + j][:k] for j in range(pruned_len - 1)]
|
||||
)
|
||||
output_top_logprobs_val.append(
|
||||
list(
|
||||
zip(
|
||||
values[pt + pruned_len - 1][:k],
|
||||
indices[pt + pruned_len - 1][:k],
|
||||
)
|
||||
values[pt + pruned_len - 1][:k],
|
||||
)
|
||||
)
|
||||
output_top_logprobs_idx.append(
|
||||
list(
|
||||
indices[pt + pruned_len - 1][:k],
|
||||
)
|
||||
)
|
||||
pt += pruned_len
|
||||
|
||||
return input_top_logprobs, output_top_logprobs
|
||||
return (
|
||||
input_top_logprobs_val,
|
||||
input_top_logprobs_idx,
|
||||
output_top_logprobs_val,
|
||||
output_top_logprobs_idx,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
@@ -193,29 +207,22 @@ class LogitsProcessor(nn.Module):
|
||||
if not logits_metadata.return_logprob:
|
||||
return LogitsProcessorOutput(
|
||||
next_token_logits=last_logits,
|
||||
next_token_logprobs=None,
|
||||
normalized_prompt_logprobs=None,
|
||||
input_token_logprobs=None,
|
||||
input_top_logprobs=None,
|
||||
output_top_logprobs=None,
|
||||
)
|
||||
else:
|
||||
last_logprobs = torch.nn.functional.log_softmax(last_logits, dim=-1)
|
||||
|
||||
if logits_metadata.forward_mode.is_decode():
|
||||
if logits_metadata.return_top_logprob:
|
||||
output_top_logprobs = self.get_top_logprobs(
|
||||
last_logprobs, logits_metadata
|
||||
)[1]
|
||||
output_top_logprobs_val, output_top_logprobs_idx = (
|
||||
self.get_top_logprobs(last_logprobs, logits_metadata)[2:4]
|
||||
)
|
||||
else:
|
||||
output_top_logprobs = None
|
||||
output_top_logprobs_val = output_top_logprobs_idx = None
|
||||
return LogitsProcessorOutput(
|
||||
next_token_logits=last_logits,
|
||||
next_token_logprobs=last_logprobs,
|
||||
normalized_prompt_logprobs=None,
|
||||
input_token_logprobs=None,
|
||||
input_top_logprobs=None,
|
||||
output_top_logprobs=output_top_logprobs,
|
||||
output_top_logprobs_val=output_top_logprobs_val,
|
||||
output_top_logprobs_idx=output_top_logprobs_idx,
|
||||
)
|
||||
else:
|
||||
# Slice the requested tokens to compute logprob
|
||||
@@ -246,11 +253,16 @@ class LogitsProcessor(nn.Module):
|
||||
|
||||
# Get the logprob of top-k tokens
|
||||
if logits_metadata.return_top_logprob:
|
||||
input_top_logprobs, output_top_logprobs = self.get_top_logprobs(
|
||||
all_logprobs, logits_metadata
|
||||
)
|
||||
(
|
||||
input_top_logprobs_val,
|
||||
input_top_logprobs_idx,
|
||||
output_top_logprobs_val,
|
||||
output_top_logprobs_idx,
|
||||
) = self.get_top_logprobs(all_logprobs, logits_metadata)
|
||||
else:
|
||||
input_top_logprobs = output_top_logprobs = None
|
||||
input_top_logprobs_val = input_top_logprobs_idx = (
|
||||
output_top_logprobs_val
|
||||
) = output_top_logprobs_idx = None
|
||||
|
||||
# Compute the normalized logprobs for the requested tokens.
|
||||
# Note that we pad a zero at the end for easy batching.
|
||||
@@ -273,8 +285,10 @@ class LogitsProcessor(nn.Module):
|
||||
next_token_logprobs=last_logprobs,
|
||||
normalized_prompt_logprobs=normalized_prompt_logprobs,
|
||||
input_token_logprobs=input_token_logprobs,
|
||||
input_top_logprobs=input_top_logprobs,
|
||||
output_top_logprobs=output_top_logprobs,
|
||||
input_top_logprobs_val=input_top_logprobs_val,
|
||||
input_top_logprobs_idx=input_top_logprobs_idx,
|
||||
output_top_logprobs_val=output_top_logprobs_val,
|
||||
output_top_logprobs_idx=output_top_logprobs_idx,
|
||||
)
|
||||
|
||||
def _get_logits(
|
||||
|
||||
Reference in New Issue
Block a user