fix(task): make detect_task architecture-aware for multi-task model types

src/winml/modelkit/loader/task.py

@@ -337,11 +337,21 @@ def detect_task(config: PretrainedConfig) -> tuple[str, str]:
     task: str | None = None
     source = "none"
 
-    # 1. Explicit (model_type, task) mapping wins.
-    for mt, mapped_task in HF_MODEL_CLASS_MAPPING:
-        if mt == model_type_normalized:
-            task, source = mapped_task, "HF_MODEL_CLASS_MAPPING"
-            break
+    # 1. Explicit (model_type, task) mapping wins — but only when unambiguous.
+    #    The mapping is keyed by (model_type, task) and a model_type may appear
+    #    under several tasks: encoder-decoder types (bart/t5: feature-extraction +
+    #    text2text-generation) plus an optional (model_type, None) default-class
+    #    sentinel. Detection therefore short-circuits only when the model_type maps
+    #    to exactly one *real* (non-None) task. With multiple distinct tasks the
+    #    architecture head is what disambiguates, so fall through to step 3 instead
+    #    of guessing; the None sentinel alone never decides the task.
+    distinct_tasks: set[str] = {
+        mapped
+        for mt, mapped in HF_MODEL_CLASS_MAPPING
+        if mt == model_type_normalized and mapped is not None
+    }
+    if len(distinct_tasks) == 1:
+        task, source = next(iter(distinct_tasks)), "HF_MODEL_CLASS_MAPPING"
 
     # 2. Wrapped-library model types (e.g. timm via "timm_wrapper") carry no
     #    `architectures`; resolve through their wrapped library instead of the

tests/integration/test_task_consistency.py

@@ -20,12 +20,26 @@
 from winml.modelkit.loader import detect_task
 
 
+# Every test here calls AutoConfig.from_pretrained, which hits the network unless
+# the HF config is already cached; mark the module so offline runs can skip via
+# ``-m 'not network'`` instead of failing with a confusing ConnectionError.
+pytestmark = pytest.mark.network
+
+
 # (model_id, expected WinMLTask) — exercises the D2 modality upgrade (DINOv2),
 # a top-level-nested vision config that must NOT upgrade (CLIP), and a text control.
 CASES = [
     ("facebook/dinov2-small", "image-feature-extraction"),
     ("openai/clip-vit-base-patch32", "feature-extraction"),  # image_size nested -> no D2
     ("bert-base-uncased", "fill-mask"),
+    # bart maps to >1 task (feature-extraction + text2text-generation) in
+    # MODEL_CLASS_MAPPING; detect_task must consult the architecture head instead
+    # of short-circuiting to the first key, so a sequence-classification BART
+    # resolves to text-classification rather than the lossy feature-extraction.
+    ("facebook/bart-large-mnli", "text-classification"),
+    # sam has a (sam, None) default-class sentinel plus a single real task; the
+    # sentinel must not make detection fall through — it resolves to mask-generation.
+    ("facebook/sam-vit-base", "mask-generation"),
 ]
 
 

tests/unit/loader/test_detect_task.py

@@ -83,6 +83,52 @@ def test_detect_task_falls_back_to_hf_task_defaults() -> None:
     assert source == "HF_TASK_DEFAULTS"
 
 
+def test_detect_task_does_not_short_circuit_for_ambiguous_model_type() -> None:
+    """A model_type with >1 distinct task in MODEL_CLASS_MAPPING (e.g. bart:
+    feature-extraction + text2text-generation) cannot be disambiguated by
+    model_type alone, so detect_task must fall through to architecture-aware
+    detection instead of short-circuiting to the first key (feature-extraction)."""
+    cfg = _FakeConfig("bart", architectures=["BartForSequenceClassification"])
+    with patch(_DETECT, return_value="text-classification") as m:
+        task, source = detect_task(cfg)
+    assert task == "text-classification"
+    assert source == "TasksManager"
+    m.assert_called_once()
+
+
+def test_detect_task_uses_single_real_task_despite_none_sentinel() -> None:
+    """A model_type with a None default-class sentinel plus exactly one real task
+    (sam: (sam, None) + (sam, mask-generation)) short-circuits to that real task
+    rather than falling through on the None sentinel."""
+    cfg = _FakeConfig("sam")
+    with patch(_DETECT) as m:
+        task, source = detect_task(cfg)
+    assert (task, source) == ("mask-generation", "HF_MODEL_CLASS_MAPPING")
+    m.assert_not_called()
+
+
+def test_detect_task_falls_through_for_multi_task_model_type_sam2() -> None:
+    """sam2 maps to multiple real tasks (image-segmentation, feature-extraction,
+    image-feature-extraction, mask-generation) in MODEL_CLASS_MAPPING, so the
+    (sam2, None) sentinel cannot disambiguate and detection must fall through to
+    architecture-aware detection rather than short-circuiting."""
+    cfg = _FakeConfig("sam2")
+    with patch(_DETECT, return_value="feature-extraction") as m:
+        task, source = detect_task(cfg)
+    assert (task, source) == ("feature-extraction", "TasksManager")
+    m.assert_called_once()
+
+
+def test_detect_task_short_circuits_for_unambiguous_model_type() -> None:
+    """A model_type with a single task entry (segformer -> image-segmentation)
+    still short-circuits via MODEL_CLASS_MAPPING without consulting TasksManager."""
+    cfg = _FakeConfig("segformer")
+    with patch(_DETECT) as m:
+        task, source = detect_task(cfg)
+    assert (task, source) == ("image-segmentation", "HF_MODEL_CLASS_MAPPING")
+    m.assert_not_called()
+
+
 def test_resolve_case1_surfaces_modality_aware_task() -> None:
     """Orchestrator Case 1 (auto-detect) surfaces the D2-upgraded task; the model
     class is unchanged (resolved from the pre-upgrade Optimum task)."""