From a15fad31577e195d6b3e0678098ebf7c7d262f0a Mon Sep 17 00:00:00 2001 From: Zhipeng Wang Date: Tue, 9 Jun 2026 12:48:45 +0800 Subject: [PATCH 1/2] fix(task): make detect_task architecture-aware for multi-task model types detect_task short-circuited on the first (model_type, task) key in MODEL_CLASS_MAPPING, so every encoder-decoder type (bart/t5/marian/...) auto-detected as feature-extraction regardless of the architecture head and disagreed with the head-aware config/build path. Short-circuit only when a model_type maps to exactly one real task; otherwise fall through to architecture-aware detection. A (model_type, None) default-class sentinel no longer forces a fall-through, so single-real-task types like sam resolve to their real task (mask-generation). --- src/winml/modelkit/loader/task.py | 20 +++++++++---- tests/integration/test_task_consistency.py | 8 +++++ tests/unit/loader/test_detect_task.py | 34 ++++++++++++++++++++++ 3 files changed, 57 insertions(+), 5 deletions(-) diff --git a/src/winml/modelkit/loader/task.py b/src/winml/modelkit/loader/task.py index 0b0187d50..392b2074f 100644 --- a/src/winml/modelkit/loader/task.py +++ b/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 = { + 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 diff --git a/tests/integration/test_task_consistency.py b/tests/integration/test_task_consistency.py index dfa260a74..2f4cb2e6c 100644 --- a/tests/integration/test_task_consistency.py +++ b/tests/integration/test_task_consistency.py @@ -26,6 +26,14 @@ ("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"), ] diff --git a/tests/unit/loader/test_detect_task.py b/tests/unit/loader/test_detect_task.py index 2f6e267f5..05c5bec09 100644 --- a/tests/unit/loader/test_detect_task.py +++ b/tests/unit/loader/test_detect_task.py @@ -83,6 +83,40 @@ 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_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).""" From 8744d1d11ca42dc89497cc20d0710520fdfefdda Mon Sep 17 00:00:00 2001 From: Zhipeng Wang Date: Tue, 9 Jun 2026 14:37:37 +0800 Subject: [PATCH 2/2] test(task): pin sam2 fall-through; add type annotation + network marker Addresses PR review: annotate distinct_tasks as set[str] (so next(iter(...)) types as str), add a sam2 multi-task fall-through regression test, and mark test_task_consistency as network since it calls AutoConfig.from_pretrained. --- src/winml/modelkit/loader/task.py | 2 +- tests/integration/test_task_consistency.py | 6 ++++++ tests/unit/loader/test_detect_task.py | 12 ++++++++++++ 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/src/winml/modelkit/loader/task.py b/src/winml/modelkit/loader/task.py index 392b2074f..36e9266d7 100644 --- a/src/winml/modelkit/loader/task.py +++ b/src/winml/modelkit/loader/task.py @@ -345,7 +345,7 @@ def detect_task(config: PretrainedConfig) -> tuple[str, str]: # 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 = { + distinct_tasks: set[str] = { mapped for mt, mapped in HF_MODEL_CLASS_MAPPING if mt == model_type_normalized and mapped is not None diff --git a/tests/integration/test_task_consistency.py b/tests/integration/test_task_consistency.py index 2f4cb2e6c..7bff0ddd5 100644 --- a/tests/integration/test_task_consistency.py +++ b/tests/integration/test_task_consistency.py @@ -20,6 +20,12 @@ 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 = [ diff --git a/tests/unit/loader/test_detect_task.py b/tests/unit/loader/test_detect_task.py index 05c5bec09..55aa0e6d1 100644 --- a/tests/unit/loader/test_detect_task.py +++ b/tests/unit/loader/test_detect_task.py @@ -107,6 +107,18 @@ def test_detect_task_uses_single_real_task_despite_none_sentinel() -> None: 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."""