tasks.py 937 Bytes
Newer Older
1
2
3
4
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Literal, get_args

5
GenerationTask = Literal["generate", "transcription", "realtime"]
6
GENERATION_TASKS: tuple[GenerationTask, ...] = get_args(GenerationTask)
7

8
PoolingTask = Literal[
9
10
11
12
13
14
15
    "embed",
    "classify",
    "score",
    "token_embed",
    "token_classify",
    "plugin",
    "embed&token_classify",
16
]
17
POOLING_TASKS: tuple[PoolingTask, ...] = get_args(PoolingTask)
18

19
20
21
22
23
24
# Score API handles score/rerank for:
# - "score" task (score_type: cross-encoder models)
# - "embed" task (score_type: bi-encoder models)
# - "token_embed" task (score_type: late interaction models)
ScoreType = Literal["bi-encoder", "cross-encoder", "late-interaction"]

25
26
27
28
FrontendTask = Literal["render"]
FRONTEND_TASKS: tuple[FrontendTask, ...] = get_args(FrontendTask)

SupportedTask = Literal[GenerationTask, PoolingTask, FrontendTask]