| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from abc import ABC, abstractmethod |
| 4 | from collections.abc import Set |
| 5 | |
| 6 | import torch |
| 7 | import torch.nn as nn |
| 8 | |
| 9 | from vllm.tasks import PoolingTask |
| 10 | from vllm.v1.outputs import PoolerOutput |
| 11 | from vllm.v1.pool.metadata import PoolingMetadata |
| 12 | |
| 13 | from .common import PoolingParamsUpdate |
| 14 | |
| 15 | |
| 16 | class Pooler(nn.Module, ABC): |
| 17 | """The interface required for all poolers used in pooling models in vLLM.""" |
| 18 | |
| 19 | @abstractmethod |
| 20 | def get_supported_tasks(self) -> Set[PoolingTask]: |
| 21 | """Determine which pooling tasks are supported.""" |
| 22 | raise NotImplementedError |
| 23 | |
| 24 | def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: |
| 25 | """Construct the updated pooling parameters to use for a supported task.""" |
| 26 | return PoolingParamsUpdate() |
| 27 | |
| 28 | @abstractmethod |
| 29 | def forward( |
| 30 | self, |
| 31 | hidden_states: torch.Tensor, |
| 32 | pooling_metadata: PoolingMetadata, |
| 33 | ) -> PoolerOutput: |
| 34 | raise NotImplementedError |
| 35 | |
| 36 | |
| 37 | __all__ = ["Pooler"] |