230 lines
8.5 KiB
Python
230 lines
8.5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from models import HFModelSummary, ModelInspection
|
|
|
|
|
|
VLLM_LIKE_FRAMEWORKS = (
|
|
"vllm",
|
|
"sglang",
|
|
"vllm-customized",
|
|
"vllm-mlu",
|
|
"vllm-016",
|
|
"vllm_fix_tokenizer",
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TaskSpec:
|
|
task_type: str
|
|
modality: str
|
|
pipeline_tags: tuple[str, ...]
|
|
priority: int
|
|
|
|
|
|
TASK_SPECS: tuple[TaskSpec, ...] = (
|
|
TaskSpec("text-generation", "text", ("text-generation",), 10),
|
|
TaskSpec(
|
|
"visual-multi-modal",
|
|
"multimodal",
|
|
("image-text-to-text", "visual-question-answering", "document-question-answering", "video-text-to-text"),
|
|
20,
|
|
),
|
|
TaskSpec("text-to-image-generation", "image", ("text-to-image", "image-to-image"), 30),
|
|
TaskSpec("asr", "audio", ("automatic-speech-recognition",), 40),
|
|
TaskSpec("question_answering", "text", ("question-answering",), 50),
|
|
TaskSpec("feature_emb", "embedding", ("feature-extraction", "sentence-similarity"), 60),
|
|
TaskSpec("vision_classification", "vision", ("image-classification", "zero-shot-image-classification"), 70),
|
|
TaskSpec("text_classification", "text", ("text-classification", "zero-shot-classification"), 80),
|
|
TaskSpec("reinforcement_learning", "text", ("reinforcement-learning",), 90),
|
|
)
|
|
|
|
|
|
TASK_SPEC_BY_TYPE = {task.task_type: task for task in TASK_SPECS}
|
|
DYNAMIC_TASK_TYPES: list[str] = []
|
|
|
|
|
|
TASK_TYPE_BY_MODEL_TASK_LEVEL_ID = {
|
|
"0": "text-generation",
|
|
"2": "visual-multi-modal",
|
|
"21": "text-to-image-generation",
|
|
"23": "text-generation",
|
|
"29": "feature_emb",
|
|
"34": "feature_emb",
|
|
"39": "question_answering",
|
|
"54": "asr",
|
|
"116": "vision_classification",
|
|
"198": "vision_classification",
|
|
"207": "reinforcement_learning",
|
|
"213": "text-to-image-generation",
|
|
}
|
|
|
|
TASK_TYPE_BY_MODEL_TASK_LEVEL = {
|
|
"文本生成": "text-generation",
|
|
"视觉多模态理解": "visual-multi-modal",
|
|
"图片生成图片": "text-to-image-generation",
|
|
"文生图": "text-to-image-generation",
|
|
"特征抽取": "feature_emb",
|
|
"句子相似度": "feature_emb",
|
|
"问答": "question_answering",
|
|
"语音识别": "asr",
|
|
"视觉分类": "vision_classification",
|
|
"零样本图像分类": "vision_classification",
|
|
"强化学习": "reinforcement_learning",
|
|
}
|
|
|
|
|
|
def task_type_from_history_task(task: dict) -> str | None:
|
|
"""Recover the API task type from current or legacy history fields."""
|
|
task_type = task.get("taskType")
|
|
if task_type and str(task_type) in TASK_SPEC_BY_TYPE:
|
|
return str(task_type)
|
|
level_id = task.get("modelTaskLevelId")
|
|
if level_id is not None:
|
|
mapped = TASK_TYPE_BY_MODEL_TASK_LEVEL_ID.get(str(level_id))
|
|
if mapped:
|
|
return mapped
|
|
level_name = task.get("modelTaskLevel")
|
|
if level_name:
|
|
return TASK_TYPE_BY_MODEL_TASK_LEVEL.get(str(level_name))
|
|
return None
|
|
|
|
|
|
def all_task_types() -> list[str]:
|
|
return [task.task_type for task in TASK_SPECS] + list(DYNAMIC_TASK_TYPES)
|
|
|
|
|
|
def register_dynamic_task_types(task_types: list[str]) -> list[str]:
|
|
"""Register API task identifiers that can be sourced by the same ModelScope tag."""
|
|
added: list[str] = []
|
|
next_priority = max(spec.priority for spec in TASK_SPEC_BY_TYPE.values()) + 10
|
|
for raw in task_types:
|
|
task_type = str(raw or "").strip()
|
|
if not task_type or task_type in TASK_SPEC_BY_TYPE:
|
|
continue
|
|
if any(not (character.isascii() and (character.isalnum() or character in "_-")) for character in task_type):
|
|
continue
|
|
TASK_SPEC_BY_TYPE[task_type] = TaskSpec(
|
|
task_type=task_type,
|
|
modality="generic",
|
|
pipeline_tags=(task_type,),
|
|
priority=next_priority,
|
|
)
|
|
next_priority += 10
|
|
DYNAMIC_TASK_TYPES.append(task_type)
|
|
added.append(task_type)
|
|
return added
|
|
|
|
|
|
def register_dynamic_task_route(task_type: str, pipeline_tag: str) -> None:
|
|
task_type = str(task_type or "").strip()
|
|
pipeline_tag = str(pipeline_tag or "").strip().lower()
|
|
if not task_type or not pipeline_tag:
|
|
return
|
|
existing = TASK_SPEC_BY_TYPE.get(task_type)
|
|
if existing is not None:
|
|
if pipeline_tag not in existing.pipeline_tags:
|
|
TASK_SPEC_BY_TYPE[task_type] = TaskSpec(
|
|
task_type=existing.task_type,
|
|
modality=existing.modality,
|
|
pipeline_tags=(*existing.pipeline_tags, pipeline_tag),
|
|
priority=existing.priority,
|
|
)
|
|
return
|
|
register_dynamic_task_types([task_type])
|
|
created = TASK_SPEC_BY_TYPE.get(task_type)
|
|
if created is not None:
|
|
TASK_SPEC_BY_TYPE[task_type] = TaskSpec(
|
|
task_type=created.task_type,
|
|
modality=created.modality,
|
|
pipeline_tags=(pipeline_tag,),
|
|
priority=created.priority,
|
|
)
|
|
|
|
|
|
def pipeline_tags_for_task_types(task_types: list[str]) -> list[str]:
|
|
tags: list[str] = []
|
|
for task_type in task_types:
|
|
spec = TASK_SPEC_BY_TYPE[task_type]
|
|
for tag in spec.pipeline_tags:
|
|
if tag not in tags:
|
|
tags.append(tag)
|
|
return tags
|
|
|
|
|
|
def task_specs_for_model(model: HFModelSummary) -> list[TaskSpec]:
|
|
pipeline_tag = (model.pipeline_tag or "").strip().lower()
|
|
return [task for task in TASK_SPEC_BY_TYPE.values() if pipeline_tag in task.pipeline_tags]
|
|
|
|
|
|
def compatible_text_generation_frameworks(
|
|
target_gpu: str,
|
|
supported_frameworks: set[str],
|
|
inspection: ModelInspection,
|
|
) -> list[str]:
|
|
can_llamacpp = "llamacpp" in supported_frameworks
|
|
vllm_like = [framework for framework in VLLM_LIKE_FRAMEWORKS if framework in supported_frameworks]
|
|
can_transformers = "transformers" in supported_frameworks
|
|
compatible: list[str] = []
|
|
|
|
if can_llamacpp and inspection.has_gguf:
|
|
compatible.append("llamacpp")
|
|
if inspection.has_vllm_weights:
|
|
compatible.extend(vllm_like)
|
|
if can_transformers:
|
|
compatible.append("transformers")
|
|
if not compatible:
|
|
raise ValueError(f"No compatible LLM weights/framework combination is available for {target_gpu}")
|
|
return compatible
|
|
|
|
|
|
def choose_text_generation_framework(target_gpu: str, supported_frameworks: set[str], inspection: ModelInspection) -> str:
|
|
return compatible_text_generation_frameworks(target_gpu, supported_frameworks, inspection)[0]
|
|
|
|
|
|
def compatible_frameworks_for_task(
|
|
task_type: str,
|
|
target_gpu: str,
|
|
supported_frameworks: set[str],
|
|
inspection: ModelInspection,
|
|
) -> list[str]:
|
|
if task_type in {"text-generation", "visual-multi-modal", "reinforcement_learning"}:
|
|
return compatible_text_generation_frameworks(target_gpu, supported_frameworks, inspection)
|
|
|
|
if task_type == "asr":
|
|
compatible: list[str] = []
|
|
if "sherpa-onnx" in supported_frameworks and inspection.has_onnx_weights:
|
|
compatible.append("sherpa-onnx")
|
|
if inspection.has_standard_weights:
|
|
compatible.extend(framework for framework in ("transformers", "funasr") if framework in supported_frameworks)
|
|
if compatible:
|
|
return compatible
|
|
raise ValueError(f"No compatible ASR framework found for {target_gpu}")
|
|
|
|
if task_type == "feature_emb":
|
|
if inspection.has_standard_weights:
|
|
compatible = [framework for framework in ("sentence-transformers", "transformers") if framework in supported_frameworks]
|
|
if compatible:
|
|
return compatible
|
|
raise ValueError(f"No compatible embedding framework found for {target_gpu}")
|
|
|
|
if task_type in {"question_answering", "vision_classification", "text_classification"}:
|
|
if "transformers" in supported_frameworks and inspection.has_standard_weights:
|
|
return ["transformers"]
|
|
raise ValueError(f"No compatible transformers template found for {task_type} on {target_gpu}")
|
|
|
|
if task_type == "text-to-image-generation":
|
|
if "diffusers" in supported_frameworks and inspection.has_standard_weights:
|
|
return ["diffusers"]
|
|
raise ValueError(f"No compatible diffusers template found for {target_gpu}")
|
|
|
|
if inspection.has_standard_weights and supported_frameworks:
|
|
return sorted(supported_frameworks)
|
|
|
|
raise ValueError(f"Unsupported task type for auto framework selection: {task_type}")
|
|
|
|
|
|
def choose_framework_for_task(task_type: str, target_gpu: str, supported_frameworks: set[str], inspection: ModelInspection) -> str:
|
|
return compatible_frameworks_for_task(task_type, target_gpu, supported_frameworks, inspection)[0]
|