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]