diff --git a/src/winml/modelkit/loader/resolution.py b/src/winml/modelkit/loader/resolution.py index 871c7c817..8f9425d80 100644 --- a/src/winml/modelkit/loader/resolution.py +++ b/src/winml/modelkit/loader/resolution.py @@ -24,6 +24,7 @@ from .task import ( HF_TASK_DEFAULTS, + KNOWN_TASKS, get_default_task_for_model_id, get_supported_tasks, normalize_task, @@ -253,6 +254,7 @@ class TaskSource(str, Enum): SENTINEL_DEFAULT = "sentinel-default" # (model_type, None) sentinel TASKS_MANAGER = "tasks-manager" # Optimum inference (incl. fill-mask upgrade) WRAPPED_LIBRARY = "wrapped-library" # no architectures -> first supported task + PIPELINE_TAG = "pipeline-tag" # Hub pipeline_tag fallback HF_TASK_DEFAULT = "hf-task-default" # last-resort default @@ -440,8 +442,8 @@ def resolve_task( """Resolve a single model's task + class from an HF config. Stages: 0 user override -> 1 detect (override / no-architectures / - TasksManager / default) -> 2 model class -> 3 modality upgrade - (detection path only) -> 4 composite tag. + TasksManager / pipeline-tag / default) -> 2 model class -> 3 modality + upgrade (detection path only) -> 4 composite tag. ``model_type_override`` lets a caller drive resolution with a build variant (e.g. ``qwen3_transformer_only``) without mutating the loaded HF config; when @@ -572,6 +574,17 @@ def resolve_task( except ValueError: opt_task = None + # 1e. Hub pipeline_tag fallback + if opt_task is None and model_id: + from ..utils.hub_utils import get_pipeline_tag + + tag = get_pipeline_tag(model_id) + if tag: + normalized_tag = normalize_task(tag) + if normalized_tag in KNOWN_TASKS: + opt_task = normalized_tag + source = TaskSource.PIPELINE_TAG + # 1d. last-resort default if opt_task is None: opt_task = next(iter(HF_TASK_DEFAULTS)) diff --git a/src/winml/modelkit/utils/hub_utils.py b/src/winml/modelkit/utils/hub_utils.py index d45adf87f..a62876d98 100644 --- a/src/winml/modelkit/utils/hub_utils.py +++ b/src/winml/modelkit/utils/hub_utils.py @@ -144,6 +144,28 @@ def is_hub_model(model_name_or_path: str) -> tuple[bool, dict]: return False, {"type": "local", "path": model_name_or_path} +_PIPELINE_TAG_TIMEOUT = 10 # seconds + + +def get_pipeline_tag(model_id: str) -> str | None: + """Return the Hub ``pipeline_tag`` for *model_id*, or ``None``. + + Lightweight helper that skips the full metadata extraction of + ``is_hub_model``. Returns ``None`` (never raises) when *model_id* is a + local path, the Hub is unreachable, or the model has no tag. + """ + if _is_local_path(model_id): + return None + try: + from huggingface_hub import HfApi + + info = HfApi().model_info(model_id, timeout=_PIPELINE_TAG_TIMEOUT) + return getattr(info, "pipeline_tag", None) + except Exception: + logger.debug("pipeline_tag lookup failed for '%s'", model_id, exc_info=True) + return None + + def inject_hub_metadata(onnx_model: Any, model_name_or_path: str, metadata: dict) -> None: """Inject HuggingFace Hub metadata into ONNX model. diff --git a/tests/unit/loader/test_detect_task.py b/tests/unit/loader/test_detect_task.py index e2a39aea6..8b8745488 100644 --- a/tests/unit/loader/test_detect_task.py +++ b/tests/unit/loader/test_detect_task.py @@ -228,3 +228,58 @@ def test_resolve_task_case1_surfaces_modality_aware_task() -> None: r = resolve_task(cfg) assert r.task == "image-feature-extraction" assert r.optimum_task == "feature-extraction" + + +# ============================================================================= +# Stage 1e — Hub pipeline_tag fallback +# ============================================================================= + +_GET_PIPELINE_TAG = "winml.modelkit.utils.hub_utils.get_pipeline_tag" + + +def test_resolve_task_uses_pipeline_tag_when_architecture_fails() -> None: + """When config.architectures contains a generic name (e.g. 'Model') that + TasksManager cannot resolve, the Hub pipeline_tag is used as fallback.""" + cfg = _FakeConfig("faketype", name_or_path="audeering/wav2vec2-large-robust-24-ft-age-gender") + with ( + patch(_INFER, side_effect=ValueError("unknown arch")), + patch(_GET_PIPELINE_TAG, return_value="audio-classification"), + ): + r = resolve_task(cfg) + assert r.task == "audio-classification" + assert r.source == TaskSource.PIPELINE_TAG + + +def test_resolve_task_pipeline_tag_skips_invalid_task() -> None: + """When pipeline_tag is not a recognized task, falls through to last-resort default.""" + cfg = _FakeConfig("faketype", name_or_path="someone/some-model") + with ( + patch(_INFER, side_effect=ValueError("unknown arch")), + patch(_GET_PIPELINE_TAG, return_value="not-a-real-task"), + ): + r = resolve_task(cfg) + assert r.source == TaskSource.HF_TASK_DEFAULT + + +def test_resolve_task_pipeline_tag_handles_api_failure() -> None: + """When the Hub API call fails (network error, etc.), get_pipeline_tag returns None.""" + cfg = _FakeConfig("faketype", name_or_path="someone/some-model") + with ( + patch(_INFER, side_effect=ValueError("unknown arch")), + patch(_GET_PIPELINE_TAG, return_value=None), + ): + r = resolve_task(cfg) + assert r.source == TaskSource.HF_TASK_DEFAULT + + +def test_resolve_task_pipeline_tag_skips_local_path() -> None: + """When model_id is a local path, get_pipeline_tag is still called but returns None + (local-path rejection is internal to get_pipeline_tag).""" + cfg = _FakeConfig("faketype", name_or_path="./local-model") + with ( + patch(_INFER, side_effect=ValueError("unknown arch")), + patch(_GET_PIPELINE_TAG, return_value=None) as mock_tag, + ): + r = resolve_task(cfg) + mock_tag.assert_called_once_with("./local-model") + assert r.source == TaskSource.HF_TASK_DEFAULT