Skip to content

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

Merged
timenick merged 6 commits into
mainfrom
zhiwang/detect-task-head-aware
Jun 9, 2026
Merged

fix(task): make detect_task architecture-aware for multi-task model types#841
timenick merged 6 commits into
mainfrom
zhiwang/detect-task-head-aware

Conversation

@timenick

@timenick timenick commented Jun 9, 2026

Copy link
Copy Markdown
Collaborator

What

detect_task short-circuited to the first (model_type, task) key in MODEL_CLASS_MAPPING, so every encoder-decoder model type (bart, t5, marian, mu2, qwen3, blip, vision-encoder-decoder) auto-detected as feature-extraction regardless of the architecture head — and disagreed with the already-head-aware config / build path (which uses _detect_task_and_class_from_config).

Change

Short-circuit step 1 only when a model_type maps to exactly one real (non-None) task. With multiple distinct tasks, 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).

Result — inspect now agrees with config / build

model before (inspect) after (inspect) config / build
facebook/bart-large-mnli feature-extraction text-classification text-classification
facebook/bart-large-cnn feature-extraction fill-mask* fill-mask*
facebook/bart-base feature-extraction feature-extraction feature-extraction
facebook/sam-vit-base feature-extraction mask-generation mask-generation
clip / siglip feature-extraction feature-extraction feature-extraction

* fill-mask for seq2seq generation models is still semantically wrong (optimum maps BartForConditionalGeneration to fill-mask). This PR only makes inspect consistent with config / build. Correcting *ForConditionalGeneration -> text2text-generation — which changes build artifacts and needs a text2text-generation export path — is Step 2, tracked in #838.

Tests

  • Unit: ambiguous multi-task type falls through to architecture-aware detection; a single-real-task type with a None sentinel (sam) resolves to its real task; a single-entry type still short-circuits.
  • Integration (test_task_consistency): added bart-large-mnli -> text-classification and sam-vit-base -> mask-generation.

Addresses #838 (Step 1 of 2).

timenick added 2 commits June 9, 2026 12:48
…ypes

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).
@timenick timenick requested a review from a team as a code owner June 9, 2026 04:50

@DingmaomaoBJTU DingmaomaoBJTU left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall this is a correct and clean fix. The core logic is sound: collect only the non-None task entries for a given model_type, and short-circuit only when there is exactly one — otherwise fall through to the architecture-aware TasksManager path. Three issues worth addressing below.

Comment thread src/winml/modelkit/loader/task.py
Comment thread src/winml/modelkit/loader/task.py
Comment thread tests/unit/loader/test_detect_task.py
Comment thread tests/integration/test_task_consistency.py
timenick added 3 commits June 9, 2026 14:37
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.

@DingmaomaoBJTU DingmaomaoBJTU left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

3 of 4 previously raised issues are resolved. One remains open (see inline comment).

  • Type safety (distinct_tasks: set[str]) ✅ Fixed.
  • sam2 unit test (test_detect_task_falls_through_for_multi_task_model_type_sam2) ✅ Added.
  • Network marker (pytestmark = pytest.mark.network) ✅ Applied module-wide with a helpful comment.
  • Performance (O(n) scan) — Still open; see inline comment. Not a blocker given the current mapping size, but worth tracking.

Comment thread src/winml/modelkit/loader/task.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants