[feat] Optional base64 encoding for the flat prompt top logprob arrays (#31960)
This commit is contained in:
@@ -211,6 +211,8 @@ class GenerateReqInput:
|
||||
# Return prompt top logprobs as flat arrays plus shape metadata instead of
|
||||
# the nested per-position [logprob, token_id, text] lists.
|
||||
return_flat_raw_top_logprobs: bool = False
|
||||
# Base64-encode the flat arrays. Requires return_flat_raw_top_logprobs.
|
||||
return_flat_raw_top_logprobs_b64: bool = False
|
||||
# Whether to stream output.
|
||||
stream: bool = False
|
||||
# Whether to log metrics for this request (e.g. health_generate calls do not log metrics)
|
||||
@@ -381,6 +383,13 @@ class GenerateReqInput:
|
||||
"scoring: delimiter-sparse top logprob rows have no contiguous "
|
||||
"position mapping."
|
||||
)
|
||||
if (
|
||||
self.return_flat_raw_top_logprobs_b64
|
||||
and not self.return_flat_raw_top_logprobs
|
||||
):
|
||||
raise ValueError(
|
||||
"return_flat_raw_top_logprobs_b64 requires return_flat_raw_top_logprobs."
|
||||
)
|
||||
|
||||
def _determine_batch_size(self):
|
||||
"""Determine if this is a single example or a batch and the batch size."""
|
||||
@@ -737,6 +746,7 @@ class GenerateReqInput:
|
||||
return_sampling_mask=self.return_sampling_mask[i],
|
||||
return_text_in_logprobs=self.return_text_in_logprobs,
|
||||
return_flat_raw_top_logprobs=self.return_flat_raw_top_logprobs,
|
||||
return_flat_raw_top_logprobs_b64=self.return_flat_raw_top_logprobs_b64,
|
||||
stream=self.stream,
|
||||
log_metrics=self.log_metrics,
|
||||
return_hidden_states=(
|
||||
|
||||
@@ -37,6 +37,7 @@ from http import HTTPStatus
|
||||
from typing import Any, Awaitable, Dict, Iterable, List, Optional, Tuple, Union
|
||||
|
||||
import fastapi
|
||||
import numpy as np
|
||||
import pybase64
|
||||
import torch
|
||||
import uvloop
|
||||
@@ -263,6 +264,7 @@ def _build_flat_input_top_logprobs_fields(
|
||||
input_top_logprobs_val: List[Optional[List[float]]],
|
||||
input_top_logprobs_idx: List[Optional[List[int]]],
|
||||
top_logprobs_num: int,
|
||||
return_b64: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
"""Build the flat raw prompt top logprob response fields.
|
||||
|
||||
@@ -270,7 +272,9 @@ def _build_flat_input_top_logprobs_fields(
|
||||
arrays. The leading null positions (counted by
|
||||
`input_top_logprobs_null_prefix`) precede the arrays, so covered position
|
||||
i, entry j lives at flat[(i - null_prefix) * k + j] and the covered range
|
||||
spans null_prefix + rows positions.
|
||||
spans null_prefix + rows positions. With ``return_b64``, the arrays are
|
||||
base64 contiguous little-endian binary; the dtype marker fields let the
|
||||
widths change later without a wire break.
|
||||
"""
|
||||
num_rows = len(input_top_logprobs_val)
|
||||
null_prefix = 0
|
||||
@@ -289,8 +293,20 @@ def _build_flat_input_top_logprobs_fields(
|
||||
)
|
||||
|
||||
fields: Dict[str, Any] = {}
|
||||
fields["input_top_logprobs_val_flat"] = [v for row in val_rows for v in row]
|
||||
fields["input_top_logprobs_idx_flat"] = [i for row in idx_rows for i in row]
|
||||
if return_b64:
|
||||
val_arr = np.asarray(val_rows, dtype=np.float32)
|
||||
idx_arr = np.asarray(idx_rows, dtype=np.int32)
|
||||
fields["input_top_logprobs_val_flat_b64"] = pybase64.b64encode(
|
||||
val_arr.tobytes()
|
||||
).decode("utf-8")
|
||||
fields["input_top_logprobs_idx_flat_b64"] = pybase64.b64encode(
|
||||
idx_arr.tobytes()
|
||||
).decode("utf-8")
|
||||
fields["input_top_logprobs_val_flat_b64_dtype"] = "float32"
|
||||
fields["input_top_logprobs_idx_flat_b64_dtype"] = "int32"
|
||||
else:
|
||||
fields["input_top_logprobs_val_flat"] = [v for row in val_rows for v in row]
|
||||
fields["input_top_logprobs_idx_flat"] = [i for row in idx_rows for i in row]
|
||||
fields["input_top_logprobs_shape"] = [len(val_rows), k]
|
||||
fields["input_top_logprobs_null_prefix"] = null_prefix
|
||||
return fields
|
||||
@@ -2301,6 +2317,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
||||
state.input_top_logprobs_val,
|
||||
state.input_top_logprobs_idx,
|
||||
top_logprobs_num,
|
||||
return_b64=state.obj.return_flat_raw_top_logprobs_b64,
|
||||
)
|
||||
)
|
||||
except ValueError as e:
|
||||
|
||||
Reference in New Issue
Block a user