main
py 45 lines 1.03 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Poolers that produce an output aggregating all tokens in the sequence."""
4
5 from .heads import (
6 ClassifierPoolerHead,
7 EmbeddingPoolerHead,
8 SequencePoolerHead,
9 SequencePoolerHeadOutput,
10 )
11 from .methods import (
12 CLSPool,
13 LastPool,
14 MeanPool,
15 SequencePoolingMethod,
16 SequencePoolingMethodOutput,
17 get_seq_pooling_method,
18 )
19 from .poolers import (
20 SequencePooler,
21 SequencePoolerOutput,
22 SequencePoolingFn,
23 SequencePoolingHeadFn,
24 pooler_for_classify,
25 pooler_for_embed,
26 )
27
28 __all__ = [
29 "SequencePoolerHead",
30 "SequencePoolerHeadOutput",
31 "ClassifierPoolerHead",
32 "EmbeddingPoolerHead",
33 "SequencePoolingMethod",
34 "SequencePoolingMethodOutput",
35 "CLSPool",
36 "LastPool",
37 "MeanPool",
38 "get_seq_pooling_method",
39 "SequencePooler",
40 "SequencePoolingFn",
41 "SequencePoolingHeadFn",
42 "SequencePoolerOutput",
43 "pooler_for_classify",
44 "pooler_for_embed",
45 ]