vllm.v1.pool.metadata

PoolingMetadata dataclass

Tensors for pooling.

Source code in vllm/v1/pool/metadata.py

@dataclass
class PoolingMetadata:
    """Tensors for pooling."""

    prompt_lens: torch.Tensor
    prompt_token_ids: Optional[torch.Tensor]
    pooling_params: list[PoolingParams]

    def __getitem__(self, indices: slice):
        return PoolingMetadata(
            prompt_lens=self.prompt_lens[indices],
            prompt_token_ids=None if self.prompt_token_ids is None else
            self.prompt_token_ids[indices],
            pooling_params=self.pooling_params[indices],
        )

pooling_params instance-attribute

pooling_params: list[PoolingParams]

prompt_lens instance-attribute

prompt_lens: Tensor

prompt_token_ids instance-attribute

prompt_token_ids: Optional[Tensor]

__getitem__

__getitem__(indices: slice)

Source code in vllm/v1/pool/metadata.py

def __getitem__(self, indices: slice):
    return PoolingMetadata(
        prompt_lens=self.prompt_lens[indices],
        prompt_token_ids=None if self.prompt_token_ids is None else
        self.prompt_token_ids[indices],
        pooling_params=self.pooling_params[indices],
    )

__init__

__init__(
    prompt_lens: Tensor,
    prompt_token_ids: Optional[Tensor],
    pooling_params: list[PoolingParams],
) -> None