adapt for ModelHub agent platform

2026-07-10 00:22:50 +08:00
commit 8c4bbff413
28 changed files with 3991 additions and 0 deletions
--- a/modelhub_submmit_api/task_registry.py
+++ b/modelhub_submmit_api/task_registry.py
@@ -0,0 +1,120 @@
+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}
+
+
+def all_task_types() -> list[str]:
+    return [task.task_type for task in TASK_SPECS]
+
+
+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_SPECS if pipeline_tag in task.pipeline_tags]
+
+
+def choose_text_generation_framework(target_gpu: str, supported_frameworks: set[str], inspection: ModelInspection) -> 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
+
+    if can_llamacpp and vllm_like:
+        if inspection.has_gguf:
+            return "llamacpp"
+        if inspection.has_vllm_weights:
+            return vllm_like[0]
+        raise ValueError(f"{target_gpu} supports both llamacpp and vLLM-like frameworks, but the model lacks usable weights")
+    if can_llamacpp:
+        if inspection.has_gguf:
+            return "llamacpp"
+        raise ValueError(f"{target_gpu} only supports llamacpp for this workflow, but no usable GGUF was found")
+    if vllm_like:
+        if inspection.has_vllm_weights:
+            return vllm_like[0]
+        raise ValueError(f"{target_gpu} only supports vLLM-like frameworks for this workflow, but no standard weights were found")
+    if can_transformers:
+        if inspection.has_vllm_weights:
+            return "transformers"
+        raise ValueError(f"{target_gpu} only supports transformers for this workflow, but no standard weights were found")
+    raise ValueError(f"No supported LLM framework template is available for {target_gpu}")
+
+
+def choose_framework_for_task(task_type: str, target_gpu: str, supported_frameworks: set[str], inspection: ModelInspection) -> str:
+    if task_type in {"text-generation", "visual-multi-modal", "reinforcement_learning"}:
+        return choose_text_generation_framework(target_gpu, supported_frameworks, inspection)
+
+    if task_type == "asr":
+        if "sherpa-onnx" in supported_frameworks and inspection.has_onnx_weights:
+            return "sherpa-onnx"
+        for framework in ("transformers", "funasr"):
+            if framework in supported_frameworks and inspection.has_standard_weights:
+                return framework
+        raise ValueError(f"No compatible ASR framework found for {target_gpu}")
+
+    if task_type == "feature_emb":
+        for framework in ("sentence-transformers", "transformers"):
+            if framework in supported_frameworks and inspection.has_standard_weights:
+                return framework
+        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}")
+
+    raise ValueError(f"Unsupported task type for auto framework selection: {task_type}")