main
py 43 lines 1.01 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Poolers that produce an output for each token in the sequence."""
4
5 from .heads import (
6 TokenClassifierPoolerHead,
7 TokenEmbeddingPoolerHead,
8 TokenPoolerHead,
9 TokenPoolerHeadOutputItem,
10 )
11 from .methods import (
12 AllPool,
13 StepPool,
14 TokenPoolingMethod,
15 TokenPoolingMethodOutputItem,
16 get_tok_pooling_method,
17 )
18 from .poolers import (
19 TokenPooler,
20 TokenPoolerOutput,
21 TokenPoolingFn,
22 TokenPoolingHeadFn,
23 pooler_for_token_classify,
24 pooler_for_token_embed,
25 )
26
27 __all__ = [
28 "TokenPoolerHead",
29 "TokenPoolerHeadOutputItem",
30 "TokenClassifierPoolerHead",
31 "TokenEmbeddingPoolerHead",
32 "TokenPoolingMethod",
33 "TokenPoolingMethodOutputItem",
34 "AllPool",
35 "StepPool",
36 "get_tok_pooling_method",
37 "TokenPooler",
38 "TokenPoolerOutput",
39 "TokenPoolingFn",
40 "TokenPoolingHeadFn",
41 "pooler_for_token_classify",
42 "pooler_for_token_embed",
43 ]