Skip to content

vllm.model_executor.layers.pooler.tokwise.poolers

TokenPoolerOutput module-attribute 

TokenPoolerOutput: TypeAlias = list[Tensor | None]

TokenPoolingFn module-attribute 

TokenPoolingFn: TypeAlias = Callable[
    [Tensor, PoolingMetadata],
    list[TokenPoolingMethodOutputItem],
]

TokenPoolingHeadFn module-attribute 

TokenPoolingHeadFn: TypeAlias = Callable[
    [list[TokenPoolingMethodOutputItem], PoolingMetadata],
    list[TokenPoolerHeadOutputItem],
]

TokenPooler

Bases: Pooler

A layer that pools specific information from hidden states.

This layer does the following: 1. Extracts specific tokens or aggregates data based on pooling method. 2. Postprocesses the output based on pooling head. 3. Returns structured results as PoolerOutput.

Source code in vllm/model_executor/layers/pooler/tokwise/poolers.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
class TokenPooler(Pooler):
    """
    A layer that pools specific information from hidden states.

    This layer does the following:
    1. Extracts specific tokens or aggregates data based on pooling method.
    2. Postprocesses the output based on pooling head.
    3. Returns structured results as `PoolerOutput`.
    """

    def __init__(
        self,
        pooling: TokenPoolingMethod | TokenPoolingFn,
        head: TokenPoolerHead | TokenPoolingHeadFn,
    ) -> None:
        super().__init__()

        self.pooling = pooling
        self.head = head

    def get_supported_tasks(self) -> Set[PoolingTask]:
        tasks = set(POOLING_TASKS)

        if isinstance(self.pooling, TokenPoolingMethod):
            tasks &= self.pooling.get_supported_tasks()
        if isinstance(self.head, TokenPoolerHead):
            tasks &= self.head.get_supported_tasks()

        return tasks

    def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
        updates = PoolingParamsUpdate()

        if isinstance(self.pooling, TokenPoolingMethod):
            updates |= self.pooling.get_pooling_updates(task)

        return updates

    def forward(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> TokenPoolerOutput:
        pooled_data = self.pooling(hidden_states, pooling_metadata)
        pooled_data = self.head(pooled_data, pooling_metadata)
        return pooled_data

head instance-attribute 

head = head

pooling instance-attribute 

pooling = pooling

__init__ 

__init__(
    pooling: TokenPoolingMethod | TokenPoolingFn,
    head: TokenPoolerHead | TokenPoolingHeadFn,
) -> None
Source code in vllm/model_executor/layers/pooler/tokwise/poolers.py 
54
55
56
57
58
59
60
61
62
def __init__(
    self,
    pooling: TokenPoolingMethod | TokenPoolingFn,
    head: TokenPoolerHead | TokenPoolingHeadFn,
) -> None:
    super().__init__()

    self.pooling = pooling
    self.head = head

forward 

forward(
    hidden_states: Tensor, pooling_metadata: PoolingMetadata
) -> TokenPoolerOutput
Source code in vllm/model_executor/layers/pooler/tokwise/poolers.py 
82
83
84
85
86
87
88
89
def forward(
    self,
    hidden_states: torch.Tensor,
    pooling_metadata: PoolingMetadata,
) -> TokenPoolerOutput:
    pooled_data = self.pooling(hidden_states, pooling_metadata)
    pooled_data = self.head(pooled_data, pooling_metadata)
    return pooled_data

get_pooling_updates 

get_pooling_updates(
    task: PoolingTask,
) -> PoolingParamsUpdate
Source code in vllm/model_executor/layers/pooler/tokwise/poolers.py 
74
75
76
77
78
79
80
def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
    updates = PoolingParamsUpdate()

    if isinstance(self.pooling, TokenPoolingMethod):
        updates |= self.pooling.get_pooling_updates(task)

    return updates

get_supported_tasks 

get_supported_tasks() -> Set[PoolingTask]
Source code in vllm/model_executor/layers/pooler/tokwise/poolers.py 
64
65
66
67
68
69
70
71
72
def get_supported_tasks(self) -> Set[PoolingTask]:
    tasks = set(POOLING_TASKS)

    if isinstance(self.pooling, TokenPoolingMethod):
        tasks &= self.pooling.get_supported_tasks()
    if isinstance(self.head, TokenPoolerHead):
        tasks &= self.head.get_supported_tasks()

    return tasks

pooler_for_token_classify 

pooler_for_token_classify(
    pooler_config: PoolerConfig,
    *,
    pooling: TokenPoolingMethod
    | TokenPoolingFn
    | None = None,
    classifier: ClassifierFn | None = None,
    act_fn: PoolerActivation | str | None = None,
)
Source code in vllm/model_executor/layers/pooler/tokwise/poolers.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def pooler_for_token_classify(
    pooler_config: PoolerConfig,
    *,
    pooling: TokenPoolingMethod | TokenPoolingFn | None = None,
    classifier: ClassifierFn | None = None,
    act_fn: PoolerActivation | str | None = None,
):
    if pooling is None:
        pooling = get_tok_pooling_method(pooler_config.get_tok_pooling_type())

    vllm_config = get_current_vllm_config()
    model_config = vllm_config.model_config
    head = TokenClassifierPoolerHead(
        head_dtype=model_config.head_dtype,
        classifier=classifier,
        logit_bias=model_config.pooler_config.logit_bias,
        activation=resolve_classifier_act_fn(
            model_config, static_num_labels=False, act_fn=act_fn
        ),
    )

    return TokenPooler(pooling=pooling, head=head)

pooler_for_token_embed 

pooler_for_token_embed(pooler_config: PoolerConfig)
Source code in vllm/model_executor/layers/pooler/tokwise/poolers.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def pooler_for_token_embed(pooler_config: PoolerConfig):
    pooling = get_tok_pooling_method(pooler_config.get_tok_pooling_type())

    vllm_config = get_current_vllm_config()
    model_config = vllm_config.model_config
    head = TokenEmbeddingPoolerHead(
        head_dtype=model_config.head_dtype,
        projector=_load_st_projector(model_config),
        activation=PoolerNormalize(),
    )

    return TokenPooler(pooling=pooling, head=head)