Fix model-task inconsistency for vision feature-extraction models

[zhenchaoni]

Fix model-task inconsistency for vision feature-extraction models

Fixes #777, #778, #782.

Principle

winml inspect is the source of truth for valid (model_id, task) pairs. Both feature-extraction and image-feature-extraction are valid ways to address an image-embedding model like facebook/dinov2-base. Downstream commands must accept whichever name winml inspect accepts, then use (model_id, task) to locate the concrete class to act on.

Root cause

Optimum's TasksManager.get_exporter_config_constructor only knows canonical Optimum task names. Several call sites passed the raw user-supplied task straight through, so HF aliases like image-feature-extraction were rejected with "Unsupported". The evaluator additionally needs to know which HF pipeline name to dispatch on, which the canonical Optimum task name doesn't carry by itself for bimodal tasks like feature-extraction.

Fix

• Inspect / export / HTP exporter: normalize via _map_task_synonym(task) (in export/io.py) before any TasksManager lookup because it requires normalized task input. This is a single function reused at each TasksManager boundary — no new global table.
• Quantize: _resolve_dataset_class(task, io_config) in datasets/__init__.py dispatches to TextDataset / ImageDataset based on the actual ONNX input names. No AutoConfig.from_pretrained round-trip. Bimodal io_configs fall back to RandomDataset with a warning.
• Evaluate: Because HF pipeline and evaluate library have their task name convention, to_hf_pipeline_task(task, model_id) in eval/evaluate.py translates to the HF pipeline name the underlying evaluate library expects. Uses OnnxConfig.inputs (no weights loaded) to pick the modality. Bimodal models (e.g. CLIP combined: both pixel_values and input_ids) keep the task unchanged via a len(hits) == 1 guard, preserving the explicit user task.

Validation

facebook/dinov2-base:

Command	Before	After
winml inspect -m facebook/dinov2-base --task image-feature-extraction	"Unsupported"	Resolves via Dinov2OnnxConfig
winml export -m facebook/dinov2-base -t image-feature-extraction	KeyError on TasksManager	Valid ONNX with last_hidden_state
winml eval -m facebook/dinov2-base --task feature-extraction	RuntimeError: Failed to create feature-extraction dataset	kNN metrics on mini-imagenet
winml quantize <onnx> --task feature-extraction -m facebook/dinov2-small	Failure by using TextDataset	Routes to ImageDataset

openai/clip-vit-base-patch32 (bimodal, regression check):

• winml eval -m openai/clip-vit-base-patch32 --task feature-extraction → stays feature-extraction (text STS evaluator); not silently rerouted to image.
• winml eval -m openai/clip-vit-base-patch32 (auto-detect) → resolves to feature-extraction (text).

Tests

Unit:

• tests/unit/eval/test_eval.py::TestResolveTask — auto-detect, explicit task, bimodal guard, HF pipeline translation.
• test_random_dataset.py — TASK_DATASET_MAPPING covers all registered tasks, including bimodal dict-of-dict.

E2E (-m e2e, dinov2 chosen because it isn't in MODEL_BUILD_CONFIGS and so actually exercises the TasksManager path):

• tests/e2e/test_inspect_e2e.py::TestInspectDinoV2 — both image-feature-extraction and feature-extraction resolve.
• tests/e2e/test_export_e2e.py::TestExportDinoV2::test_image_feature_extraction.
• tests/e2e/test_eval_e2e.py::TestEvalPerTask::test_image_feature_extraction parameterized over both task names.
• tests/e2e/test_quantize_e2e.py::test_feature_extraction_with_pixel_values_uses_image_dataset.

[timenick]

Thanks for tackling these three — the e2e coverage on dinov2 end-to-end across inspect/export/eval/quantize is exactly what was missing. Two concrete concerns and a proposal before this lands.

Concrete regression: CLIP feature-extraction will swap dataset class

canonical_task_to_known_task iterates _OPTIMUM_CANONICAL_TO_KNOWN_TASK["feature-extraction"] and returns on first input-name match. For CLIP, CLIPOnnxConfig for feature-extraction declares both pixel_values and input_ids (verified via TasksManager.get_exporter_config_constructor("onnx", model_type="clip", task="feature-extraction") — inputs are ['pixel_values', 'input_ids', 'attention_mask']). With current dict insertion order, pixel_values matches first → routes to image-feature-extraction.

But src/winml/modelkit/models/hf/clip.py:53 already encodes the right answer:

("clip", "feature-extraction"): CLIPTextModelWithProjection,        # text encoder
("clip", "image-feature-extraction"): CLIPVisionModelWithProjection, # vision encoder

So winml quantize -m openai/clip-vit-base-patch32 --task feature-extraction:

• Before this PR: loader picks CLIPTextModelWithProjection, datasets picks TextDataset → consistent
• After this PR: loader still picks CLIPTextModelWithProjection, datasets translated to image-feature-extraction → ImageDataset → dtype/shape mismatch at calibration

The input-name heuristic doesn't have enough information to disambiguate multi-modal architectures; MODEL_CLASS_MAPPING already holds the truth and disagrees with what the heuristic returns.

DatasetCalibrationReader.__init__ quietly rewrites self.task

Caller passes task="feature-extraction", but after the new lines self.task becomes "image-feature-extraction". Any downstream consumer of reader.task expecting "what the caller asked for" gets something different. Either name the new attribute (self.resolved_task?) or document the semantic shift explicitly on the class.

Architectural direction — two patches vs. one source fix

The pattern we're ending up with: three places inline-call _map_task_synonym (forward), plus a recovery function (reverse) that has to re-derive what was lost. Both directions exist because _detect_task_from_config drops modality at the source by calling TasksManager.infer_task_from_model, which is Optimum-canonical.

Prevention alternative: have _detect_task_from_config return the HF (modality-aware) task to begin with (HF Hub pipeline_tag first, then a (model_type, config)→HF modality fallback). Then:

• the three inline forward-direction calls collapse into one boundary helper (we'd want this regardless — a lint rule against direct TasksManager imports keeps new code honest),
• canonical_task_to_known_task and _OPTIMUM_CANONICAL_TO_KNOWN_TASK go away,
• the CLIP-style ambiguity above doesn't exist because we never need to recover modality from input names.

Proposed path

1. Fix the CLIP regression in this PR — it's a marquee architecture and the regression is real.
2. Resolve the self.task semantic shift (rename or document).
3. Merge — the e2e fix value for Model task handling is not consistent between config and build #777/winml build and eval command fails on image embedding model like facebook/dinov2-base #778/winml inspect -m facebook/dinov2-base --task image-feature-extraction returns unsupported #782 is real and shouldn't wait.
4. I'll open a follow-up doing the prevention refactor: _detect_task_from_config → HF task, centralize the forward-normalize into a boundary helper with a lint rule, delete the canonical-to-known table. Your e2e tests stay green throughout.

Happy to pair on the follow-up — there's overlap with the #724 / KNOWN_TASKS cleanup I'm queuing up separately.

---

(Edit: removed accidental #1/#2/#3 section markers that GitHub auto-linked to unrelated issues.)

[DingmaomaoBJTU]

Good fix for the task-naming inconsistency between HF pipeline aliases and Optimum's canonical names. The approach of normalizing via _map_task_synonym at each TasksManager call site is correct. The to_hf_pipeline_task evaluator mapper and the bimodal feature-extraction dataset dispatch are both well-thought-out. A few observations: