diff --git a/examples/recipes/dandelin_vilt-b32-finetuned-vqa/cpu/cpu/visual-question-answering_fp16_config.json b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/cpu/cpu/visual-question-answering_fp16_config.json new file mode 100644 index 000000000..0a8ae56a3 --- /dev/null +++ b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/cpu/cpu/visual-question-answering_fp16_config.json @@ -0,0 +1 @@ +{"export":{"opset_version":17,"batch_size":1,"export_params":true,"do_constant_folding":true,"verbose":false,"dynamo":false,"enable_hierarchy_tags":true,"clean_onnx":false,"hierarchy_tag_format":"full","input_tensors":[{"name":"input_ids","dtype":"int32","shape":[1,40],"value_range":[0,30522]},{"name":"attention_mask","dtype":"int32","shape":[1,40],"value_range":[1,2]},{"name":"token_type_ids","dtype":"int32","shape":[1,40],"value_range":[0,2]},{"name":"pixel_values","dtype":"float32","shape":[1,3,384,384],"value_range":[0,1]}],"output_tensors":[{"name":"logits"}]},"optim":{},"quant":{"mode":"fp16","samples":10,"calibration_method":"minmax","weight_type":"uint8","activation_type":"uint8","per_channel":false,"symmetric":false,"fp16_keep_io_types":true},"compile":null,"loader":{"task":"visual-question-answering","model_class":"ViltForQuestionAnswering","model_type":"vilt"}} diff --git a/examples/recipes/dandelin_vilt-b32-finetuned-vqa/cpu/cpu/visual-question-answering_fp32_config.json b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/cpu/cpu/visual-question-answering_fp32_config.json new file mode 100644 index 000000000..a8ea92c3a --- /dev/null +++ b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/cpu/cpu/visual-question-answering_fp32_config.json @@ -0,0 +1 @@ +{"export":{"opset_version":17,"batch_size":1,"export_params":true,"do_constant_folding":true,"verbose":false,"dynamo":false,"enable_hierarchy_tags":true,"clean_onnx":false,"hierarchy_tag_format":"full","input_tensors":[{"name":"input_ids","dtype":"int32","shape":[1,40],"value_range":[0,30522]},{"name":"attention_mask","dtype":"int32","shape":[1,40],"value_range":[1,2]},{"name":"token_type_ids","dtype":"int32","shape":[1,40],"value_range":[0,2]},{"name":"pixel_values","dtype":"float32","shape":[1,3,384,384],"value_range":[0,1]}],"output_tensors":[{"name":"logits"}]},"optim":{},"quant":null,"compile":null,"loader":{"task":"visual-question-answering","model_class":"ViltForQuestionAnswering","model_type":"vilt"}} diff --git a/examples/recipes/dandelin_vilt-b32-finetuned-vqa/dml/gpu/visual-question-answering_fp16_config.json b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/dml/gpu/visual-question-answering_fp16_config.json new file mode 100644 index 000000000..0a8ae56a3 --- /dev/null +++ b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/dml/gpu/visual-question-answering_fp16_config.json @@ -0,0 +1 @@ +{"export":{"opset_version":17,"batch_size":1,"export_params":true,"do_constant_folding":true,"verbose":false,"dynamo":false,"enable_hierarchy_tags":true,"clean_onnx":false,"hierarchy_tag_format":"full","input_tensors":[{"name":"input_ids","dtype":"int32","shape":[1,40],"value_range":[0,30522]},{"name":"attention_mask","dtype":"int32","shape":[1,40],"value_range":[1,2]},{"name":"token_type_ids","dtype":"int32","shape":[1,40],"value_range":[0,2]},{"name":"pixel_values","dtype":"float32","shape":[1,3,384,384],"value_range":[0,1]}],"output_tensors":[{"name":"logits"}]},"optim":{},"quant":{"mode":"fp16","samples":10,"calibration_method":"minmax","weight_type":"uint8","activation_type":"uint8","per_channel":false,"symmetric":false,"fp16_keep_io_types":true},"compile":null,"loader":{"task":"visual-question-answering","model_class":"ViltForQuestionAnswering","model_type":"vilt"}} diff --git a/examples/recipes/dandelin_vilt-b32-finetuned-vqa/dml/gpu/visual-question-answering_fp32_config.json b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/dml/gpu/visual-question-answering_fp32_config.json new file mode 100644 index 000000000..a8ea92c3a --- /dev/null +++ b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/dml/gpu/visual-question-answering_fp32_config.json @@ -0,0 +1 @@ +{"export":{"opset_version":17,"batch_size":1,"export_params":true,"do_constant_folding":true,"verbose":false,"dynamo":false,"enable_hierarchy_tags":true,"clean_onnx":false,"hierarchy_tag_format":"full","input_tensors":[{"name":"input_ids","dtype":"int32","shape":[1,40],"value_range":[0,30522]},{"name":"attention_mask","dtype":"int32","shape":[1,40],"value_range":[1,2]},{"name":"token_type_ids","dtype":"int32","shape":[1,40],"value_range":[0,2]},{"name":"pixel_values","dtype":"float32","shape":[1,3,384,384],"value_range":[0,1]}],"output_tensors":[{"name":"logits"}]},"optim":{},"quant":null,"compile":null,"loader":{"task":"visual-question-answering","model_class":"ViltForQuestionAnswering","model_type":"vilt"}} diff --git a/examples/recipes/dandelin_vilt-b32-finetuned-vqa/openvino/cpu/visual-question-answering_fp32_config.json b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/openvino/cpu/visual-question-answering_fp32_config.json new file mode 100644 index 000000000..da38df1d8 --- /dev/null +++ b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/openvino/cpu/visual-question-answering_fp32_config.json @@ -0,0 +1 @@ +{"export":{"opset_version":17,"batch_size":1,"export_params":true,"do_constant_folding":true,"verbose":false,"dynamo":false,"enable_hierarchy_tags":true,"clean_onnx":false,"hierarchy_tag_format":"full","input_tensors":[{"name":"input_ids","dtype":"int32","shape":[1,40],"value_range":[0,30522]},{"name":"attention_mask","dtype":"int32","shape":[1,40],"value_range":[1,2]},{"name":"token_type_ids","dtype":"int32","shape":[1,40],"value_range":[0,2]},{"name":"pixel_values","dtype":"float32","shape":[1,3,384,384],"value_range":[0,1]}],"output_tensors":[{"name":"logits"}]},"optim":{},"quant":null,"compile":null,"loader":{"task":"visual-question-answering","model_class":"ViltForQuestionAnswering","model_type":"vilt"}} \ No newline at end of file diff --git a/examples/recipes/dandelin_vilt-b32-finetuned-vqa/openvino/gpu/visual-question-answering_fp32_config.json b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/openvino/gpu/visual-question-answering_fp32_config.json new file mode 100644 index 000000000..da38df1d8 --- /dev/null +++ b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/openvino/gpu/visual-question-answering_fp32_config.json @@ -0,0 +1 @@ +{"export":{"opset_version":17,"batch_size":1,"export_params":true,"do_constant_folding":true,"verbose":false,"dynamo":false,"enable_hierarchy_tags":true,"clean_onnx":false,"hierarchy_tag_format":"full","input_tensors":[{"name":"input_ids","dtype":"int32","shape":[1,40],"value_range":[0,30522]},{"name":"attention_mask","dtype":"int32","shape":[1,40],"value_range":[1,2]},{"name":"token_type_ids","dtype":"int32","shape":[1,40],"value_range":[0,2]},{"name":"pixel_values","dtype":"float32","shape":[1,3,384,384],"value_range":[0,1]}],"output_tensors":[{"name":"logits"}]},"optim":{},"quant":null,"compile":null,"loader":{"task":"visual-question-answering","model_class":"ViltForQuestionAnswering","model_type":"vilt"}} \ No newline at end of file diff --git a/examples/recipes/dandelin_vilt-b32-finetuned-vqa/openvino/npu/visual-question-answering_fp32_config.json b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/openvino/npu/visual-question-answering_fp32_config.json new file mode 100644 index 000000000..da38df1d8 --- /dev/null +++ b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/openvino/npu/visual-question-answering_fp32_config.json @@ -0,0 +1 @@ +{"export":{"opset_version":17,"batch_size":1,"export_params":true,"do_constant_folding":true,"verbose":false,"dynamo":false,"enable_hierarchy_tags":true,"clean_onnx":false,"hierarchy_tag_format":"full","input_tensors":[{"name":"input_ids","dtype":"int32","shape":[1,40],"value_range":[0,30522]},{"name":"attention_mask","dtype":"int32","shape":[1,40],"value_range":[1,2]},{"name":"token_type_ids","dtype":"int32","shape":[1,40],"value_range":[0,2]},{"name":"pixel_values","dtype":"float32","shape":[1,3,384,384],"value_range":[0,1]}],"output_tensors":[{"name":"logits"}]},"optim":{},"quant":null,"compile":null,"loader":{"task":"visual-question-answering","model_class":"ViltForQuestionAnswering","model_type":"vilt"}} \ No newline at end of file diff --git a/examples/recipes/dandelin_vilt-b32-finetuned-vqa/qnn/npu/visual-question-answering_fp32_config.json b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/qnn/npu/visual-question-answering_fp32_config.json new file mode 100644 index 000000000..da38df1d8 --- /dev/null +++ b/examples/recipes/dandelin_vilt-b32-finetuned-vqa/qnn/npu/visual-question-answering_fp32_config.json @@ -0,0 +1 @@ +{"export":{"opset_version":17,"batch_size":1,"export_params":true,"do_constant_folding":true,"verbose":false,"dynamo":false,"enable_hierarchy_tags":true,"clean_onnx":false,"hierarchy_tag_format":"full","input_tensors":[{"name":"input_ids","dtype":"int32","shape":[1,40],"value_range":[0,30522]},{"name":"attention_mask","dtype":"int32","shape":[1,40],"value_range":[1,2]},{"name":"token_type_ids","dtype":"int32","shape":[1,40],"value_range":[0,2]},{"name":"pixel_values","dtype":"float32","shape":[1,3,384,384],"value_range":[0,1]}],"output_tensors":[{"name":"logits"}]},"optim":{},"quant":null,"compile":null,"loader":{"task":"visual-question-answering","model_class":"ViltForQuestionAnswering","model_type":"vilt"}} \ No newline at end of file diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index 3c5b0c12b..d085ff84c 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/src/winml/modelkit/models/hf/__init__.py @@ -92,6 +92,8 @@ from .t5 import T5_CONFIG from .t5 import T5DecoderIOConfig as _T5DecoderIOConfig # triggers registration from .t5 import T5EncoderIOConfig as _T5EncoderIOConfig # triggers registration +from .vilt import MODEL_CLASS_MAPPING as _VILT_CLASS_MAPPING +from .vilt import ViltVqaOnnxConfig as _ViltVqaOnnxConfig # triggers registration from .vision_encoder_decoder import MODEL_CLASS_MAPPING as _VED_CLASS_MAPPING from .vision_encoder_decoder import VISION_ENCODER_DECODER_CONFIG from .vision_encoder_decoder import ( @@ -131,6 +133,7 @@ _SIGLIP_CLASS_MAPPING, _T5_CLASS_MAPPING, _VED_CLASS_MAPPING, + _VILT_CLASS_MAPPING, _VITPOSE_CLASS_MAPPING, ) for _key, _model_cls in _sub_mapping.items() diff --git a/src/winml/modelkit/models/hf/vilt.py b/src/winml/modelkit/models/hf/vilt.py new file mode 100644 index 000000000..1999328d8 --- /dev/null +++ b/src/winml/modelkit/models/hf/vilt.py @@ -0,0 +1,271 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""ViLT (Vision-and-Language Transformer) HuggingFace Model Configuration. + +ViLT is a single-stream multi-modal transformer that processes text + image +in a unified attention stack. The ``ViltForQuestionAnswering`` head produces +classification logits over the answer vocabulary configured by the checkpoint. + +Optimum has NO vendor ``ViltOnnxConfig`` (verified 2026-06-24: ``vilt`` is +absent from ``TasksManager._SUPPORTED_MODEL_TYPE`` for the transformers +library). This module writes the export config from scratch. + +The forward takes 4 required tensors (pixel_mask is omitted — see Notes): + - ``pixel_values`` [B, 3, 384, 384] RGB image + - ``input_ids`` [B, 40] tokenized question + - ``attention_mask`` [B, 40] text padding mask + - ``token_type_ids`` [B, 40] BERT segment IDs (modality) + +Output: ``logits`` [B, num_labels] over the answer vocabulary. + +Notes: +----- +ViLT's stock ``visual_embed`` is fundamentally NOT ONNX-traceable: it iterates +Python-level over tensor values (``for h, w in zip(x_h, x_w)``), uses +``torch.multinomial`` (random + non-exportable), and does per-row Python loops +over ``nonzero()`` results. We replace it during export with a statically- +shaped equivalent (see ``_patched_visual_embed`` + ``_ViltVisualEmbedPatcher``) +that assumes an all-ones ``pixel_mask`` — which is what ``ViltProcessor`` emits +for a square 384x384 input (see ``inputs`` for the square-only constraint). +Because the patched path ignores ``pixel_mask``, we drop it from the exported +ONNX graph. +""" + +from __future__ import annotations + +import types +from types import TracebackType +from typing import Any, cast + +import torch +from optimum.exporters.onnx import OnnxConfig +from optimum.exporters.onnx.model_patcher import ModelPatcher +from optimum.utils import NormalizedTextConfig +from optimum.utils.input_generators import DummyVisionInputGenerator +from transformers import ViltForQuestionAnswering + +from ...export import MaxLengthTextInputGenerator, register_onnx_overwrite + + +# ============================================================================= +# Export-time patch for ``ViltEmbeddings.visual_embed`` +# ============================================================================= +# Upstream ``visual_embed`` is **not ONNX-traceable** as written: +# * ``for h, w in zip(x_h, x_w)`` iterates Python-level over tensor values +# * ``nonzero()`` + ``unique()`` + per-row Python list-comprehension subset +# selection over a dynamic ``valid_idx`` +# * ``torch.multinomial`` random sampling (non-deterministic, not exportable) +# The eager path silently "works" only when ``pixel_mask`` is all-ones and the +# batch is 1, because the Python loop runs once with a concrete value. Under +# legacy ``torch.onnx.export`` tracing the shape resolves to 0 and PyTorch's +# ``F.interpolate`` aborts with ``input (H: 12, W: 12) output (H: 0, W: 0)``. +# +# For the production ``visual-question-answering`` inference path with a square +# 384x384 image the ``ViltProcessor`` emits an all-ones ``pixel_mask``, +# so the per-sample subset selection is a no-op. We replace ``visual_embed`` +# during export with a simplified, statically-shaped implementation that: +# * uses ``x.shape[2], x.shape[3]`` (static) for position-embed interpolation +# * skips ``multinomial`` / nonzero / Python-level batch loops +# * returns an all-ones token mask of length ``H*W + 1`` (patches + CLS) +# +def _patched_visual_embed( + self: Any, + pixel_values: torch.Tensor, + pixel_mask: torch.Tensor | None, + max_image_length: int = 200, +) -> tuple[torch.Tensor, torch.Tensor, None]: + """Static-shape, ONNX-traceable replacement for ``ViltEmbeddings.visual_embed``.""" + from torch import nn + + x = self.patch_embeddings(pixel_values) + batch_size, num_channels, height, width = x.shape + + patch_dim = self.config.image_size // self.config.patch_size + spatial_pos = self.position_embeddings[:, 1:, :].transpose(1, 2).view( + 1, num_channels, patch_dim, patch_dim + ) + pos_embed = nn.functional.interpolate( + spatial_pos, size=(height, width), mode="bilinear", align_corners=True + ) + pos_embed = pos_embed.flatten(2).transpose(1, 2).expand(batch_size, -1, -1) + + x = x.flatten(2).transpose(1, 2) + + cls_tokens = self.cls_token.expand(batch_size, -1, -1) + x = torch.cat((cls_tokens, x), dim=1) + pos_cls = self.position_embeddings[:, 0:1, :].expand(batch_size, -1, -1) + pos_embed = torch.cat((pos_cls, pos_embed), dim=1) + x = x + pos_embed + x = self.dropout(x) + + num_tokens = height * width + 1 # patches + CLS + x_mask = torch.ones(batch_size, num_tokens, dtype=torch.long, device=x.device) + return x, x_mask, None + + +class _ViltVisualEmbedPatcher(ModelPatcher): # type: ignore[misc] # untyped Optimum base + """Swaps ``ViltEmbeddings.visual_embed`` for the duration of ONNX export.""" + + def __enter__(self) -> _ViltVisualEmbedPatcher: + super().__enter__() + emb = ( + self._model.vilt.embeddings + if hasattr(self._model, "vilt") + else self._model.embeddings + ) + self._emb_ref = emb + self._orig_visual_embed = emb.visual_embed + emb.visual_embed = types.MethodType(_patched_visual_embed, emb) + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self._emb_ref.visual_embed = self._orig_visual_embed + super().__exit__(exc_type, exc_value, traceback) + + +# ============================================================================= +# Optimum ONNX Export Config Registration +# ============================================================================= +@register_onnx_overwrite("vilt", "visual-question-answering", library_name="transformers") +class ViltVqaOnnxConfig(OnnxConfig): # type: ignore[misc] # untyped Optimum base + """ONNX export config for ``ViltForQuestionAnswering``. + + Declares 4 multi-modal inputs (text triple + pixel_values) and the single + classification output. ``pixel_mask`` is deliberately omitted — see + ``inputs`` property docstring and the module-level ``Notes`` section for + the full rationale. + + Inputs: + - ``input_ids``: [B, 40] int64 + - ``attention_mask``: [B, 40] int64 + - ``token_type_ids``: [B, 40] int64 + - ``pixel_values``: [B, 3, 384, 384] float32 + + Outputs: + - ``logits``: [B, num_labels] float32 + + Notes: + - ``num_labels`` is a config-time fact, not declared dynamic in the + symbolic axes — it's a static dim of ``logits``. + - ``sequence_length`` resolves to ``max_position_embeddings`` (40 for + ViLT-B/32) via ``NORMALIZED_CONFIG_CLASS``; the + ``MaxLengthTextInputGenerator`` reads this for dummy tokens. + - Chained ``DummyVisionInputGenerator`` + ``MaxLengthTextInputGenerator`` + produce ``pixel_values`` + ``input_ids``/``attention_mask``/ + ``token_type_ids``. The patched ``visual_embed`` (see module-level + ``_ViltVisualEmbedPatcher``) synthesizes an all-ones token mask + internally, so no ``pixel_mask`` input is required. + """ + + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( + sequence_length="max_position_embeddings", + num_channels="num_channels", + image_size="image_size", + patch_size="patch_size", + allow_new=True, + ) + + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyVisionInputGenerator, + MaxLengthTextInputGenerator, + ) + + DEFAULT_ONNX_OPSET = 17 + + @property + def inputs(self) -> dict[str, dict[int, str]]: + """Declare 4 model inputs (insertion order matches forward). + + ``pixel_values`` H,W is kept STATIC at ``image_size`` (384x384), so the + exported ONNX accepts ONLY 384x384 pixel_values. + + Honest constraint: ``ViltImageProcessor`` pins the *shortest* edge to + 384 with ``size_divisor=32`` and preserves aspect ratio, so ONLY square + inputs land on 384x384 — a non-square image (e.g. 480x640 -> 384x512) + does NOT match this graph and must be square-resized upstream, which + distorts aspect ratio and can cost VQA accuracy. Callers must feed + 384x384 (square) pixel_values. + + Dynamic H,W is a known follow-up, not enabled here: the patched + ``visual_embed`` already interpolates position embeddings to the real + ``x.shape[2], x.shape[3]``, so a dynamic-H,W export is plausible — but + it is left static because it has NOT been export-verified. (The original + 0x0 ``Resize`` shape-inference failure was a property of ViLT's *stock* + non-traceable ``visual_embed``, which the patcher replaces; it does not + by itself justify pinning the patched path.) + + Note: ViLT's ``forward`` also takes a ``pixel_mask`` parameter, but + this contribution exports without it. For the square-384 path the + ``ViltProcessor`` emits an all-ones mask, and our export-time + ``ModelPatcher`` replaces the original ``visual_embed`` with a + statically-shaped version that synthesizes an all-ones token mask + internally. Including ``pixel_mask`` as an ONNX input would + dead-code-eliminate (since the patched path doesn't reference it) and + confuse runtime callers. + """ + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + "token_type_ids": {0: "batch_size", 1: "sequence_length"}, + "pixel_values": {0: "batch_size"}, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + """Single classification output over fixed answer vocabulary.""" + return { + "logits": {0: "batch_size"}, + } + + def generate_dummy_inputs( + self, + framework: str = "pt", + **kwargs: Any, + ) -> dict[str, Any]: + """Generate the 4 declared inputs via the chained vendor generators. + + ``pixel_mask`` is intentionally NOT generated — see ``inputs`` docstring. + Our model patcher's replacement ``visual_embed`` synthesizes an + all-ones token mask internally, so the model can be called with the + 4 declared inputs. + """ + dummy = cast( + "dict[str, Any]", + super().generate_dummy_inputs(framework=framework, **kwargs), + ) + # Drop any pixel_mask the generators may have produced — the patched + # visual_embed ignores it (and including it would error at sess.run + # since it isn't in the exported ONNX graph). + dummy.pop("pixel_mask", None) + return dummy + + def patch_model_for_export( + self, + model: Any, + model_kwargs: dict[str, Any] | None = None, + ) -> _ViltVisualEmbedPatcher: + """Install the ``visual_embed`` patcher for the export context.""" + return _ViltVisualEmbedPatcher(self, model, model_kwargs=model_kwargs) + + +# ============================================================================= +# HuggingFace Model Class Mapping +# ============================================================================= +# ``visual-question-answering`` has no default AutoModel routing for ViLT; +# bind the (model_type, task) tuple directly to the head-bearing HF class. +MODEL_CLASS_MAPPING: dict[tuple[str, str], type] = { + ("vilt", "visual-question-answering"): ViltForQuestionAnswering, +} + + +__all__ = [ + "MODEL_CLASS_MAPPING", + "ViltVqaOnnxConfig", +] diff --git a/tests/unit/export/test_vilt_onnx_config.py b/tests/unit/export/test_vilt_onnx_config.py new file mode 100644 index 000000000..33e63f953 --- /dev/null +++ b/tests/unit/export/test_vilt_onnx_config.py @@ -0,0 +1,174 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Tests for ViLT VQA ONNX config registration and I/O specs. + +Verifies that ViltVqaOnnxConfig is correctly registered with Optimum's +TasksManager for visual-question-answering, and that I/O specs match the +expected tensors: + +- Inputs: input_ids, attention_mask, token_type_ids, pixel_values +- Outputs: logits (over the fixed answer vocabulary) + +``pixel_mask`` is intentionally NOT an ONNX input: the export-time +ModelPatcher replaces ViLT's non-traceable ``visual_embed`` with a +statically-shaped version that synthesizes an all-ones token mask +internally, so the exported graph is called with only the 4 declared +inputs. This module locks that contract in. + +``visual-question-answering`` has no default AutoModel routing for ViLT, +so the (model_type, task) tuple is bound directly to the head-bearing HF +class via MODEL_CLASS_MAPPING. + +See also: modelkit/models/hf/vilt.py +""" + +from __future__ import annotations + +import pytest +from optimum.exporters.tasks import TasksManager + +from winml.modelkit.export import resolve_io_specs +from winml.modelkit.models import HF_MODEL_CLASS_MAPPING +from winml.modelkit.models.hf.vilt import ViltVqaOnnxConfig + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def vilt_config(): + """Minimal ViltConfig matching the exported dandelin/vilt-b32 topology. + + image_size/patch_size/num_channels drive the vision path; the hidden + dims are shrunk to keep the (weightless) dummy-input generation fast. + """ + from transformers import ViltConfig + + return ViltConfig( + image_size=384, + patch_size=32, + num_channels=3, + max_position_embeddings=40, + hidden_size=32, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=37, + num_labels=8, + ) + + +# ============================================================================= +# Registration verification +# ============================================================================= + + +class TestViltRegistration: + """Verify ViltVqaOnnxConfig is registered for visual-question-answering.""" + + def test_vilt_config_registered(self) -> None: + """ViltVqaOnnxConfig must be registered for visual-question-answering.""" + config_constructor = TasksManager.get_exporter_config_constructor( + exporter="onnx", + model_type="vilt", + task="visual-question-answering", + library_name="transformers", + ) + actual_class_name = config_constructor.func.__name__ + assert actual_class_name == "ViltVqaOnnxConfig", ( + f"Expected ViltVqaOnnxConfig for vilt/visual-question-answering, " + f"got {actual_class_name}. Registration may have failed." + ) + + def test_aggregate_mapping_binds_vqa_head(self) -> None: + """The imported architecture mapping must reach the aggregate registry.""" + from transformers import ViltForQuestionAnswering + + assert ( + HF_MODEL_CLASS_MAPPING[("vilt", "visual-question-answering")] + is ViltForQuestionAnswering + ) + + +# ============================================================================= +# I/O verification +# ============================================================================= + + +class TestViltIOSpecs: + """Verify I/O specs for ViLT visual question answering.""" + + def test_declares_exactly_four_inputs(self, vilt_config) -> None: + """Inputs must be exactly the 4 declared text+vision tensors.""" + specs = resolve_io_specs("vilt", "visual-question-answering", vilt_config) + assert set(specs["input_names"]) == { + "input_ids", + "attention_mask", + "token_type_ids", + "pixel_values", + } + + def test_pixel_mask_is_dropped(self, vilt_config) -> None: + """pixel_mask must NOT be an exported input (patched visual_embed synthesizes it).""" + specs = resolve_io_specs("vilt", "visual-question-answering", vilt_config) + assert "pixel_mask" not in specs["input_names"], ( + "pixel_mask must not be exported: the patched visual_embed builds an " + "all-ones mask internally, and exporting it would break sess.run." + ) + + def test_outputs_contain_logits(self, vilt_config) -> None: + """Output must be the single classification logits tensor.""" + specs = resolve_io_specs("vilt", "visual-question-answering", vilt_config) + assert specs["output_names"] == ["logits"] + + def test_pixel_values_shape_is_static_384(self, vilt_config) -> None: + """pixel_values H,W must be pinned to the config image_size (static).""" + specs = resolve_io_specs("vilt", "visual-question-answering", vilt_config) + idx = specs["input_names"].index("pixel_values") + _, channels, height, width = specs["input_shapes"][idx] + assert (channels, height, width) == (3, 384, 384) + + +# ============================================================================= +# Model class mapping verification +# ============================================================================= + + +class TestViltModelClassMapping: + """Verify the (model_type, task) tuple binds to the head-bearing HF class.""" + + def test_mapping_binds_vqa_head(self) -> None: + """(vilt, visual-question-answering) must map to ViltForQuestionAnswering.""" + from transformers import ViltForQuestionAnswering + + from winml.modelkit.models.hf.vilt import MODEL_CLASS_MAPPING + + assert ( + MODEL_CLASS_MAPPING[("vilt", "visual-question-answering")] + is ViltForQuestionAnswering + ) + + +class TestViltVisualEmbedPatcher: + """Verify the export patch is installed only for the export context.""" + + def test_patcher_replaces_and_restores_visual_embed(self, vilt_config) -> None: + """The patcher must not leave a mutated model after export.""" + from transformers import ViltForQuestionAnswering + + model = ViltForQuestionAnswering(vilt_config) + onnx_config = ViltVqaOnnxConfig( + vilt_config, + task="visual-question-answering", + ) + embeddings = model.vilt.embeddings + original = embeddings.visual_embed.__func__ + + with onnx_config.patch_model_for_export(model): + assert embeddings.visual_embed.__func__.__name__ == "_patched_visual_embed" + + assert embeddings.visual_embed.__func__ is original