vllm.entrypoints.pooling.utils

Functions:

• encode_pooling_output_float_or_ndarray –

  Return an ndarray when the response renderer can serialize NumPy.

• reserve_tokens –

  Shrink every length limit by num_tokens that are appended afterwards.

encode_pooling_output_float_or_ndarray(output)

Return an ndarray when the response renderer can serialize NumPy.

Source code in vllm/entrypoints/pooling/utils.py

def encode_pooling_output_float_or_ndarray(output: PoolingRequestOutput) -> Any:
    """Return an ndarray when the response renderer can serialize NumPy."""
    try:
        data = output.outputs.data
        if not data.is_contiguous():
            data = data.contiguous()
        return data.numpy()
    except (RuntimeError, TypeError):
        return output.outputs.data.tolist()

reserve_tokens(tok_params, num_tokens)

Shrink every length limit by num_tokens that are appended afterwards.

Source code in vllm/entrypoints/pooling/utils.py

def reserve_tokens(tok_params: TokenizeParams, num_tokens: int) -> TokenizeParams:
    """Shrink every length limit by `num_tokens` that are appended afterwards."""

    def shrink(length: int | None) -> int | None:
        return length if length is None or length < 0 else max(length - num_tokens, 0)

    return replace(
        tok_params,
        max_total_tokens=shrink(tok_params.max_total_tokens),
        truncate_prompt_tokens=shrink(tok_params.truncate_prompt_tokens),
        pad_prompt_tokens=shrink(tok_params.pad_prompt_tokens),
    )