Skip to content

vllm.model_executor.layers.pooler

AllPool

Bases: SimplePooler

Source code in vllm/model_executor/layers/pooler.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
class AllPool(SimplePooler):

    def extract_states(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> Union[list[torch.Tensor], torch.Tensor]:
        prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)

        offset = 0
        pooled_data = list[torch.Tensor]()
        for prompt_len in prompt_lens:
            pooled_data.append(hidden_states[offset:offset + prompt_len])
            offset += prompt_len

        return pooled_data

extract_states 

extract_states(
    hidden_states: Tensor, pooling_metadata: PoolingMetadata
) -> Union[list[Tensor], Tensor]
Source code in vllm/model_executor/layers/pooler.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
def extract_states(
    self,
    hidden_states: torch.Tensor,
    pooling_metadata: PoolingMetadata,
) -> Union[list[torch.Tensor], torch.Tensor]:
    prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)

    offset = 0
    pooled_data = list[torch.Tensor]()
    for prompt_len in prompt_lens:
        pooled_data.append(hidden_states[offset:offset + prompt_len])
        offset += prompt_len

    return pooled_data

CLSPool

Bases: SimplePooler

Source code in vllm/model_executor/layers/pooler.py 
105
106
107
108
109
110
111
112
113
114
115
116
class CLSPool(SimplePooler):

    def extract_states(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> Union[list[torch.Tensor], torch.Tensor]:
        prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)

        first_token_flat_indices = torch.zeros_like(prompt_lens)
        first_token_flat_indices[1:] += torch.cumsum(prompt_lens, dim=0)[:-1]
        return hidden_states[first_token_flat_indices]

extract_states 

extract_states(
    hidden_states: Tensor, pooling_metadata: PoolingMetadata
) -> Union[list[Tensor], Tensor]
Source code in vllm/model_executor/layers/pooler.py 
107
108
109
110
111
112
113
114
115
116
def extract_states(
    self,
    hidden_states: torch.Tensor,
    pooling_metadata: PoolingMetadata,
) -> Union[list[torch.Tensor], torch.Tensor]:
    prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)

    first_token_flat_indices = torch.zeros_like(prompt_lens)
    first_token_flat_indices[1:] += torch.cumsum(prompt_lens, dim=0)[:-1]
    return hidden_states[first_token_flat_indices]

CrossEncodingPooler

Bases: Module

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. Normalizes output if specified. 3. Returns structured results as PoolerOutput.

Attributes:

Name Type Description
pooling_type

The type of pooling to use.
normalize

Whether to normalize the pooled data.
Source code in vllm/model_executor/layers/pooler.py 
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
class CrossEncodingPooler(nn.Module):
    """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. Normalizes output if specified.
    3. Returns structured results as `PoolerOutput`.

    Attributes:
        pooling_type: The type of pooling to use.
        normalize: Whether to normalize the pooled data.
    """

    def __init__(
        self,
        config: PretrainedConfig,
        classifier: nn.Module,
        pooler: Optional[nn.Module] = None,
    ):
        super().__init__()
        self.classifier = classifier
        self.pooler = pooler
        self.default_activation_function = \
            get_cross_encoder_activation_function(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> PoolerOutput:
        """Pools sentence pair scores from the hidden_states."""

        prompt_lens = PoolingTensors.from_pooling_metadata(
            pooling_metadata, hidden_states.device).prompt_lens

        offset = 0
        pooled_data_lst = []
        for prompt_len in prompt_lens:
            pooled_data_i = hidden_states[offset:offset + prompt_len]

            if self.pooler is not None:
                final_shape_tensor = self.pooler(pooled_data_i)
            else:
                final_shape_tensor = self.classifier(pooled_data_i)

            pooled_data_lst.append(final_shape_tensor)
            offset += prompt_len

        pooled_output = torch.stack(pooled_data_lst)

        if self.pooler is not None:
            # apply classifier once on the full batch if possible
            pooled_output = self.classifier(pooled_output)

        scores = self.default_activation_function(pooled_output).squeeze(-1)

        pooled_outputs = [PoolingSequenceGroupOutput(data) for data in scores]
        return PoolerOutput(outputs=pooled_outputs)

classifier instance-attribute 

classifier = classifier

default_activation_function instance-attribute 

default_activation_function = (
    get_cross_encoder_activation_function(config)
)

pooler instance-attribute 

pooler = pooler

__init__ 

__init__(
    config: PretrainedConfig,
    classifier: Module,
    pooler: Optional[Module] = None,
)
Source code in vllm/model_executor/layers/pooler.py 
299
300
301
302
303
304
305
306
307
308
309
def __init__(
    self,
    config: PretrainedConfig,
    classifier: nn.Module,
    pooler: Optional[nn.Module] = None,
):
    super().__init__()
    self.classifier = classifier
    self.pooler = pooler
    self.default_activation_function = \
        get_cross_encoder_activation_function(config)

forward 

forward(
    hidden_states: Tensor, pooling_metadata: PoolingMetadata
) -> PoolerOutput

Pools sentence pair scores from the hidden_states.

Source code in vllm/model_executor/layers/pooler.py 
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
def forward(
    self,
    hidden_states: torch.Tensor,
    pooling_metadata: PoolingMetadata,
) -> PoolerOutput:
    """Pools sentence pair scores from the hidden_states."""

    prompt_lens = PoolingTensors.from_pooling_metadata(
        pooling_metadata, hidden_states.device).prompt_lens

    offset = 0
    pooled_data_lst = []
    for prompt_len in prompt_lens:
        pooled_data_i = hidden_states[offset:offset + prompt_len]

        if self.pooler is not None:
            final_shape_tensor = self.pooler(pooled_data_i)
        else:
            final_shape_tensor = self.classifier(pooled_data_i)

        pooled_data_lst.append(final_shape_tensor)
        offset += prompt_len

    pooled_output = torch.stack(pooled_data_lst)

    if self.pooler is not None:
        # apply classifier once on the full batch if possible
        pooled_output = self.classifier(pooled_output)

    scores = self.default_activation_function(pooled_output).squeeze(-1)

    pooled_outputs = [PoolingSequenceGroupOutput(data) for data in scores]
    return PoolerOutput(outputs=pooled_outputs)

LastPool

Bases: SimplePooler

Source code in vllm/model_executor/layers/pooler.py 
119
120
121
122
123
124
125
126
127
128
129
class LastPool(SimplePooler):

    def extract_states(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> Union[list[torch.Tensor], torch.Tensor]:
        prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)

        last_token_flat_indices = torch.cumsum(prompt_lens, dim=0) - 1
        return hidden_states[last_token_flat_indices]

extract_states 

extract_states(
    hidden_states: Tensor, pooling_metadata: PoolingMetadata
) -> Union[list[Tensor], Tensor]
Source code in vllm/model_executor/layers/pooler.py 
121
122
123
124
125
126
127
128
129
def extract_states(
    self,
    hidden_states: torch.Tensor,
    pooling_metadata: PoolingMetadata,
) -> Union[list[torch.Tensor], torch.Tensor]:
    prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)

    last_token_flat_indices = torch.cumsum(prompt_lens, dim=0) - 1
    return hidden_states[last_token_flat_indices]

MeanPool

Bases: SimplePooler

Source code in vllm/model_executor/layers/pooler.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
class MeanPool(SimplePooler):

    def extract_states(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> Union[list[torch.Tensor], torch.Tensor]:
        prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)

        cumsum = torch.cumsum(hidden_states, dim=0)
        start_indices = torch.cat([
            torch.tensor([0], device=hidden_states.device),
            torch.cumsum(prompt_lens[:-1], dim=0)
        ])
        end_indices = torch.cumsum(prompt_lens, dim=0)
        return (cumsum[end_indices - 1] - cumsum[start_indices] +
                hidden_states[start_indices]) / prompt_lens.unsqueeze(1)

extract_states 

extract_states(
    hidden_states: Tensor, pooling_metadata: PoolingMetadata
) -> Union[list[Tensor], Tensor]
Source code in vllm/model_executor/layers/pooler.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
def extract_states(
    self,
    hidden_states: torch.Tensor,
    pooling_metadata: PoolingMetadata,
) -> Union[list[torch.Tensor], torch.Tensor]:
    prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)

    cumsum = torch.cumsum(hidden_states, dim=0)
    start_indices = torch.cat([
        torch.tensor([0], device=hidden_states.device),
        torch.cumsum(prompt_lens[:-1], dim=0)
    ])
    end_indices = torch.cumsum(prompt_lens, dim=0)
    return (cumsum[end_indices - 1] - cumsum[start_indices] +
            hidden_states[start_indices]) / prompt_lens.unsqueeze(1)

Pooler

Bases: Module

Source code in vllm/model_executor/layers/pooler.py 
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
class Pooler(nn.Module):

    @classmethod
    def from_config_with_defaults(
        cls,
        pooler_config: PoolerConfig,
        pooling_type: PoolingType,
        normalize: bool,
        softmax: bool,
        step_tag_id: Optional[int] = None,
        returned_token_ids: Optional[list[int]] = None,
    ) -> SimplePooler:
        return SimplePooler.from_pooling_type(
            pooling_type=PoolingType[pooler_config.pooling_type]
            if pooler_config.pooling_type is not None else pooling_type,
            normalize=pooler_config.normalize
            if pooler_config.normalize is not None else normalize,
            softmax=pooler_config.softmax
            if pooler_config.softmax is not None else softmax,
            step_tag_id=pooler_config.step_tag_id
            if pooler_config.step_tag_id is not None else step_tag_id,
            returned_token_ids=pooler_config.returned_token_ids
            if pooler_config.returned_token_ids is not None else
            returned_token_ids,
        )

from_config_with_defaults classmethod 

from_config_with_defaults(
    pooler_config: PoolerConfig,
    pooling_type: PoolingType,
    normalize: bool,
    softmax: bool,
    step_tag_id: Optional[int] = None,
    returned_token_ids: Optional[list[int]] = None,
) -> SimplePooler
Source code in vllm/model_executor/layers/pooler.py 
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
@classmethod
def from_config_with_defaults(
    cls,
    pooler_config: PoolerConfig,
    pooling_type: PoolingType,
    normalize: bool,
    softmax: bool,
    step_tag_id: Optional[int] = None,
    returned_token_ids: Optional[list[int]] = None,
) -> SimplePooler:
    return SimplePooler.from_pooling_type(
        pooling_type=PoolingType[pooler_config.pooling_type]
        if pooler_config.pooling_type is not None else pooling_type,
        normalize=pooler_config.normalize
        if pooler_config.normalize is not None else normalize,
        softmax=pooler_config.softmax
        if pooler_config.softmax is not None else softmax,
        step_tag_id=pooler_config.step_tag_id
        if pooler_config.step_tag_id is not None else step_tag_id,
        returned_token_ids=pooler_config.returned_token_ids
        if pooler_config.returned_token_ids is not None else
        returned_token_ids,
    )

PoolerHead

Bases: Module

Source code in vllm/model_executor/layers/pooler.py 
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
class PoolerHead(nn.Module):

    def __init__(self, *, normalize: bool, softmax: bool) -> None:
        super().__init__()

        self.normalize = normalize
        self.softmax = softmax

    def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
                pooling_metadata: PoolingMetadata):

        dimensions_list = [
            pooling_param.dimensions
            for _, pooling_param in pooling_metadata.seq_groups
        ]
        if any(d is not None for d in dimensions_list):
            # change the output dimension
            assert len(pooled_data) == len(dimensions_list)
            pooled_data = [
                vecs if d is None else vecs[..., :d]
                for vecs, d in zip(pooled_data, dimensions_list)
            ]

        if self.normalize:
            if isinstance(pooled_data, list):
                pooled_data = [
                    F.normalize(data, p=2, dim=-1) for data in pooled_data
                ]
            else:
                pooled_data = F.normalize(pooled_data, p=2, dim=-1)

        if self.softmax:
            if isinstance(pooled_data, list):
                pooled_data = [
                    F.softmax(data, dim=-1)
                    if data.shape[-1] >= 2 else F.sigmoid(data)
                    for data in pooled_data
                ]
            else:
                if pooled_data.shape[-1] >= 2:
                    pooled_data = F.softmax(pooled_data, dim=-1)
                else:
                    pooled_data = F.sigmoid(pooled_data)

        return pooled_data

normalize instance-attribute 

normalize = normalize

softmax instance-attribute 

softmax = softmax

__init__ 

__init__(*, normalize: bool, softmax: bool) -> None
Source code in vllm/model_executor/layers/pooler.py 
214
215
216
217
218
def __init__(self, *, normalize: bool, softmax: bool) -> None:
    super().__init__()

    self.normalize = normalize
    self.softmax = softmax

forward 

forward(
    pooled_data: Union[list[Tensor], Tensor],
    pooling_metadata: PoolingMetadata,
)
Source code in vllm/model_executor/layers/pooler.py 
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
            pooling_metadata: PoolingMetadata):

    dimensions_list = [
        pooling_param.dimensions
        for _, pooling_param in pooling_metadata.seq_groups
    ]
    if any(d is not None for d in dimensions_list):
        # change the output dimension
        assert len(pooled_data) == len(dimensions_list)
        pooled_data = [
            vecs if d is None else vecs[..., :d]
            for vecs, d in zip(pooled_data, dimensions_list)
        ]

    if self.normalize:
        if isinstance(pooled_data, list):
            pooled_data = [
                F.normalize(data, p=2, dim=-1) for data in pooled_data
            ]
        else:
            pooled_data = F.normalize(pooled_data, p=2, dim=-1)

    if self.softmax:
        if isinstance(pooled_data, list):
            pooled_data = [
                F.softmax(data, dim=-1)
                if data.shape[-1] >= 2 else F.sigmoid(data)
                for data in pooled_data
            ]
        else:
            if pooled_data.shape[-1] >= 2:
                pooled_data = F.softmax(pooled_data, dim=-1)
            else:
                pooled_data = F.sigmoid(pooled_data)

    return pooled_data

PoolingType

Bases: IntEnum

Enumeration for different types of pooling methods.

Source code in vllm/model_executor/layers/pooler.py 
20
21
22
23
24
25
26
class PoolingType(IntEnum):
    """Enumeration for different types of pooling methods."""
    LAST = 0
    ALL = 1
    CLS = 2
    STEP = 3
    MEAN = 4

ALL class-attribute instance-attribute 

ALL = 1

CLS class-attribute instance-attribute 

CLS = 2

LAST class-attribute instance-attribute 

LAST = 0

MEAN class-attribute instance-attribute 

MEAN = 4

STEP class-attribute instance-attribute 

STEP = 3

SimplePooler

Bases: Module

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. Normalizes output if specified. 3. Returns structured results as PoolerOutput.

Attributes:

Name Type Description
pooling_type

The type of pooling to use.
normalize

Whether to normalize the pooled data.
Source code in vllm/model_executor/layers/pooler.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 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
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
class SimplePooler(nn.Module):
    """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. Normalizes output if specified.
    3. Returns structured results as `PoolerOutput`.

    Attributes:
        pooling_type: The type of pooling to use.
        normalize: Whether to normalize the pooled data.
    """

    @staticmethod
    def from_pooling_type(
        pooling_type: PoolingType,
        *,
        normalize: bool,
        softmax: bool,
        step_tag_id: Optional[int] = None,
        returned_token_ids: Optional[list[int]] = None,
    ) -> "SimplePooler":
        if pooling_type == PoolingType.LAST:
            assert step_tag_id is None and returned_token_ids is None
            return LastPool(normalize=normalize, softmax=softmax)
        if pooling_type == PoolingType.ALL:
            assert step_tag_id is None and returned_token_ids is None
            return AllPool(normalize=normalize, softmax=softmax)
        if pooling_type == PoolingType.CLS:
            assert step_tag_id is None and returned_token_ids is None
            return CLSPool(normalize=normalize, softmax=softmax)
        if pooling_type == PoolingType.MEAN:
            assert step_tag_id is None and returned_token_ids is None
            return MeanPool(normalize=normalize, softmax=softmax)
        if pooling_type == PoolingType.STEP:
            return StepPool(normalize=normalize,
                            softmax=softmax,
                            step_tag_id=step_tag_id,
                            returned_token_ids=returned_token_ids)

        assert_never(pooling_type)

    def __init__(self, *, normalize: bool, softmax: bool) -> None:
        super().__init__()

        self.head = PoolerHead(normalize=normalize, softmax=softmax)

    def get_prompt_lens(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> torch.Tensor:
        return PoolingTensors.from_pooling_metadata(
            pooling_metadata, hidden_states.device).prompt_lens

    def extract_states(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> Union[list[torch.Tensor], torch.Tensor]:
        raise NotImplementedError

    def build_output(self, data: torch.Tensor) -> PoolingSequenceGroupOutput:
        return PoolingSequenceGroupOutput(data)

    def forward(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> PoolerOutput:
        pooled_data = self.extract_states(hidden_states, pooling_metadata)
        pooled_data = self.head(pooled_data, pooling_metadata)
        pooled_outputs = [self.build_output(data) for data in pooled_data]
        return PoolerOutput(outputs=pooled_outputs)

head instance-attribute 

head = PoolerHead(normalize=normalize, softmax=softmax)

__init__ 

__init__(*, normalize: bool, softmax: bool) -> None
Source code in vllm/model_executor/layers/pooler.py 
71
72
73
74
def __init__(self, *, normalize: bool, softmax: bool) -> None:
    super().__init__()

    self.head = PoolerHead(normalize=normalize, softmax=softmax)

build_output 

build_output(data: Tensor) -> PoolingSequenceGroupOutput
Source code in vllm/model_executor/layers/pooler.py 
91
92
def build_output(self, data: torch.Tensor) -> PoolingSequenceGroupOutput:
    return PoolingSequenceGroupOutput(data)

extract_states 

extract_states(
    hidden_states: Tensor, pooling_metadata: PoolingMetadata
) -> Union[list[Tensor], Tensor]
Source code in vllm/model_executor/layers/pooler.py 
84
85
86
87
88
89
def extract_states(
    self,
    hidden_states: torch.Tensor,
    pooling_metadata: PoolingMetadata,
) -> Union[list[torch.Tensor], torch.Tensor]:
    raise NotImplementedError

forward 

forward(
    hidden_states: Tensor, pooling_metadata: PoolingMetadata
) -> PoolerOutput
Source code in vllm/model_executor/layers/pooler.py 
 94
 95
 96
 97
 98
 99
100
101
102
def forward(
    self,
    hidden_states: torch.Tensor,
    pooling_metadata: PoolingMetadata,
) -> PoolerOutput:
    pooled_data = self.extract_states(hidden_states, pooling_metadata)
    pooled_data = self.head(pooled_data, pooling_metadata)
    pooled_outputs = [self.build_output(data) for data in pooled_data]
    return PoolerOutput(outputs=pooled_outputs)

from_pooling_type staticmethod 

from_pooling_type(
    pooling_type: PoolingType,
    *,
    normalize: bool,
    softmax: bool,
    step_tag_id: Optional[int] = None,
    returned_token_ids: Optional[list[int]] = None,
) -> SimplePooler
Source code in vllm/model_executor/layers/pooler.py 
42
43
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
@staticmethod
def from_pooling_type(
    pooling_type: PoolingType,
    *,
    normalize: bool,
    softmax: bool,
    step_tag_id: Optional[int] = None,
    returned_token_ids: Optional[list[int]] = None,
) -> "SimplePooler":
    if pooling_type == PoolingType.LAST:
        assert step_tag_id is None and returned_token_ids is None
        return LastPool(normalize=normalize, softmax=softmax)
    if pooling_type == PoolingType.ALL:
        assert step_tag_id is None and returned_token_ids is None
        return AllPool(normalize=normalize, softmax=softmax)
    if pooling_type == PoolingType.CLS:
        assert step_tag_id is None and returned_token_ids is None
        return CLSPool(normalize=normalize, softmax=softmax)
    if pooling_type == PoolingType.MEAN:
        assert step_tag_id is None and returned_token_ids is None
        return MeanPool(normalize=normalize, softmax=softmax)
    if pooling_type == PoolingType.STEP:
        return StepPool(normalize=normalize,
                        softmax=softmax,
                        step_tag_id=step_tag_id,
                        returned_token_ids=returned_token_ids)

    assert_never(pooling_type)

get_prompt_lens 

get_prompt_lens(
    hidden_states: Tensor, pooling_metadata: PoolingMetadata
) -> Tensor
Source code in vllm/model_executor/layers/pooler.py 
76
77
78
79
80
81
82
def get_prompt_lens(
    self,
    hidden_states: torch.Tensor,
    pooling_metadata: PoolingMetadata,
) -> torch.Tensor:
    return PoolingTensors.from_pooling_metadata(
        pooling_metadata, hidden_states.device).prompt_lens

StepPool

Bases: SimplePooler

Source code in vllm/model_executor/layers/pooler.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
class StepPool(SimplePooler):

    def __init__(
        self,
        *,
        normalize: bool,
        softmax: bool,
        step_tag_id: Optional[int] = None,
        returned_token_ids: Optional[list[int]] = None,
    ):
        super().__init__(normalize=normalize, softmax=softmax)

        self.step_tag_id = step_tag_id
        self.returned_token_ids = returned_token_ids

    def extract_states(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> Union[list[torch.Tensor], torch.Tensor]:
        prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)

        returned_token_ids = self.returned_token_ids
        if returned_token_ids is not None and len(returned_token_ids) > 0:
            hidden_states = hidden_states[:, returned_token_ids]

        step_tag_id = self.step_tag_id

        offset = 0
        pooled_data = list[torch.Tensor]()
        for prompt_len, seq_data_i in zip(prompt_lens,
                                          pooling_metadata.seq_data.values()):
            pooled_data_i = hidden_states[offset:offset + prompt_len]
            if step_tag_id is not None:
                token_ids = torch.tensor(seq_data_i.prompt_token_ids)
                pooled_data_i = pooled_data_i[token_ids == step_tag_id]

            offset += prompt_len
            pooled_data.append(pooled_data_i)

        return pooled_data

returned_token_ids instance-attribute 

returned_token_ids = returned_token_ids

step_tag_id instance-attribute 

step_tag_id = step_tag_id

__init__ 

__init__(
    *,
    normalize: bool,
    softmax: bool,
    step_tag_id: Optional[int] = None,
    returned_token_ids: Optional[list[int]] = None,
)
Source code in vllm/model_executor/layers/pooler.py 
171
172
173
174
175
176
177
178
179
180
181
182
def __init__(
    self,
    *,
    normalize: bool,
    softmax: bool,
    step_tag_id: Optional[int] = None,
    returned_token_ids: Optional[list[int]] = None,
):
    super().__init__(normalize=normalize, softmax=softmax)

    self.step_tag_id = step_tag_id
    self.returned_token_ids = returned_token_ids

extract_states 

extract_states(
    hidden_states: Tensor, pooling_metadata: PoolingMetadata
) -> Union[list[Tensor], Tensor]
Source code in vllm/model_executor/layers/pooler.py 
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
def extract_states(
    self,
    hidden_states: torch.Tensor,
    pooling_metadata: PoolingMetadata,
) -> Union[list[torch.Tensor], torch.Tensor]:
    prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)

    returned_token_ids = self.returned_token_ids
    if returned_token_ids is not None and len(returned_token_ids) > 0:
        hidden_states = hidden_states[:, returned_token_ids]

    step_tag_id = self.step_tag_id

    offset = 0
    pooled_data = list[torch.Tensor]()
    for prompt_len, seq_data_i in zip(prompt_lens,
                                      pooling_metadata.seq_data.values()):
        pooled_data_i = hidden_states[offset:offset + prompt_len]
        if step_tag_id is not None:
            token_ids = torch.tensor(seq_data_i.prompt_token_ids)
            pooled_data_i = pooled_data_i[token_ids == step_tag_id]

        offset += prompt_len
        pooled_data.append(pooled_data_i)

    return pooled_data