40 lines
893 B
Python
40 lines
893 B
Python
from typing import Literal
|
|
|
|
|
|
class TaskType:
|
|
@classmethod
|
|
def values(cls):
|
|
"""Return a set of all valid task type values."""
|
|
return {
|
|
value
|
|
for key, value in vars(cls).items()
|
|
if not key.startswith("_") and isinstance(value, str)
|
|
}
|
|
|
|
|
|
class vLLMTaskType(TaskType):
|
|
"""The type of task to run on the vLLM engine."""
|
|
|
|
# Generate text.
|
|
GENERATE = "generate"
|
|
|
|
# Generate embeddings.
|
|
EMBED = "embed"
|
|
|
|
# Classification (e.g., sequence classification models).
|
|
CLASSIFY = "classify"
|
|
|
|
# Scoring (e.g., cross-encoder models).
|
|
SCORE = "score"
|
|
|
|
|
|
class SGLangTaskType(TaskType):
|
|
"""The type of task to run on the SGLang engine."""
|
|
|
|
# Generate text.
|
|
GENERATE = "generate"
|
|
|
|
|
|
TypeVLLMTaskType = Literal[tuple(vLLMTaskType.values())]
|
|
TypeSGLangTaskType = Literal[tuple(SGLangTaskType.values())]
|