recipe(layoutlm): add document QA support for impira/layoutlm-document-qa

examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp16_config.json

@@ -0,0 +1,81 @@
+{
+  "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,
+          512
+        ],
+        "value_range": [
+          0,
+          50265
+        ]
+      },
+      {
+        "name": "bbox",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512,
+          4
+        ],
+        "value_range": [
+          0,
+          1000
+        ]
+      },
+      {
+        "name": "attention_mask",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512
+        ],
+        "value_range": [
+          0,
+          2
+        ]
+      },
+      {
+        "name": "token_type_ids",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512
+        ],
+        "value_range": [
+          0,
+          1
+        ]
+      }
+    ],
+    "output_tensors": [
+      {
+        "name": "start_logits"
+      },
+      {
+        "name": "end_logits"
+      }
+    ]
+  },
+  "optim": {
+    "clamp_constant_values": true
+  },
+  "quant": null,
+  "loader": {
+    "task": "question-answering",
+    "model_class": "LayoutLMForQuestionAnswering",
+    "model_type": "layoutlm"
+  }
+}

examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp32_config.json

@@ -0,0 +1,81 @@
+{
+  "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,
+          512
+        ],
+        "value_range": [
+          0,
+          50265
+        ]
+      },
+      {
+        "name": "bbox",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512,
+          4
+        ],
+        "value_range": [
+          0,
+          1000
+        ]
+      },
+      {
+        "name": "attention_mask",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512
+        ],
+        "value_range": [
+          0,
+          2
+        ]
+      },
+      {
+        "name": "token_type_ids",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512
+        ],
+        "value_range": [
+          0,
+          1
+        ]
+      }
+    ],
+    "output_tensors": [
+      {
+        "name": "start_logits"
+      },
+      {
+        "name": "end_logits"
+      }
+    ]
+  },
+  "optim": {
+    "clamp_constant_values": true
+  },
+  "quant": null,
+  "loader": {
+    "task": "question-answering",
+    "model_class": "LayoutLMForQuestionAnswering",
+    "model_type": "layoutlm"
+  }
+}

src/winml/modelkit/models/hf/__init__.py

@@ -44,6 +44,7 @@
 from .depth_anything import DepthAnythingIOConfig as _DepthAnythingIOConfig  # triggers registration
 from .depth_pro import DepthProIOConfig as _DepthProIOConfig  # triggers registration
 from .detr import DETR_CONFIG
+from .layoutlm import LayoutLMQAIOConfig as _LayoutLMQAIOConfig  # triggers registration
 from .marian import MARIAN_CONFIG
 from .marian import MODEL_CLASS_MAPPING as _MARIAN_CLASS_MAPPING
 from .marian import MarianDecoderIOConfig as _MarianDecoderIOConfig  # triggers registration

src/winml/modelkit/models/hf/layoutlm.py

@@ -0,0 +1,68 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""LayoutLM HuggingFace Model Configuration."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, cast
+
+from optimum.exporters.onnx.model_configs import LayoutLMOnnxConfig
+from optimum.utils import NormalizedTextConfig
+from optimum.utils.input_generators import DummyBboxInputGenerator, DummyVisionInputGenerator
+
+from ...export import MaxLengthTextInputGenerator, register_onnx_overwrite
+
+
+if TYPE_CHECKING:
+    import torch
+
+
+class ZeroTokenTypeLayoutLMTextInputGenerator(MaxLengthTextInputGenerator):
+    """LayoutLM text dummy generator that keeps token_type_ids within type_vocab_size=1."""
+
+    def generate(
+        self,
+        input_name: str,
+        framework: str = "pt",
+        int_dtype: str = "int64",
+        float_dtype: str = "fp32",
+    ) -> torch.Tensor:
+        """Generate LayoutLM text inputs, replacing token_type_ids with zeros."""
+        tensor = cast(
+            "torch.Tensor",
+            super().generate(
+                input_name,
+                framework=framework,
+                int_dtype=int_dtype,
+                float_dtype=float_dtype,
+            ),
+        )
+        if input_name == "token_type_ids":
+            return tensor.new_zeros(tensor.shape)
+        return tensor
+
+
+@register_onnx_overwrite("layoutlm", "question-answering", library_name="transformers")
+class LayoutLMQAIOConfig(LayoutLMOnnxConfig):  # type: ignore[misc]  # optimum base is untyped
+    """LayoutLM question-answering OnnxConfig with bbox and safe token type IDs."""
+
+    # sequence_length is bound to the model's max_position_embeddings so
+    # MaxLengthTextInputGenerator emits full-length text inputs instead of
+    # Optimum's default of 16 (allow_new=True permits adding this mapping).
+    # We deliberately do NOT map max_2d_position_embeddings here: Optimum's
+    # DummyBboxInputGenerator hardcodes its coordinate range (its
+    # normalized_config.max_2d_position_embeddings read is commented out
+    # upstream), so such a mapping is inert and never becomes a sequence
+    # length. bbox coordinate bounds for the shipped recipe come from the
+    # recipe's `value_range` instead.
+    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(
+        sequence_length="max_position_embeddings",
+        allow_new=True,
+    )
+    DUMMY_INPUT_GENERATOR_CLASSES: tuple[type[Any], ...] = (
+        ZeroTokenTypeLayoutLMTextInputGenerator,
+        DummyVisionInputGenerator,
+        DummyBboxInputGenerator,
+    )

src/winml/modelkit/onnx/io.py

@@ -106,6 +106,19 @@ def to_tensor(self) -> Any:
             lo, hi = self.value_range
             if torch_dtype.is_floating_point:
                 return torch.rand(concrete_shape, dtype=torch_dtype) * (hi - lo) + lo
+            # bbox uses the HF convention: a [..., 4] tensor of box corners
+            # (x0, y0, x1, y1). We key on the input name because "bbox" is the
+            # canonical HF/Optimum name for 2D box coordinates; a model using a
+            # different name (e.g. "bboxes") simply falls through to the generic
+            # randint path below. Reorder so x0<=x1 and y0<=y1 to keep each box
+            # well-formed for layout models.
+            if self.name == "bbox" and len(concrete_shape) >= 1 and concrete_shape[-1] == 4:
+                coords = torch.randint(int(lo), int(hi), concrete_shape, dtype=torch_dtype)
+                x0 = torch.minimum(coords[..., 0], coords[..., 2])
+                y0 = torch.minimum(coords[..., 1], coords[..., 3])
+                x1 = torch.maximum(coords[..., 0], coords[..., 2])
+                y1 = torch.maximum(coords[..., 1], coords[..., 3])
+                return torch.stack((x0, y0, x1, y1), dim=-1)
             return torch.randint(int(lo), int(hi), concrete_shape, dtype=torch_dtype)
 
         # Fallback: no range info (backward compatible)

tests/unit/export/test_onnx_config_overrides.py

@@ -57,6 +57,7 @@ class TestOnnxConfigRegistration:
             ("xlm-roberta", "fill-mask", "XLMRobertaIOConfig"),
             ("camembert", "fill-mask", "CamemBERTIOConfig"),
             ("mpnet", "fill-mask", "MPNetIOConfig"),
+            ("layoutlm", "question-answering", "LayoutLMQAIOConfig"),
             ("zoedepth", "depth-estimation", "ZoeDepthIOConfig"),
         ],
         ids=[
@@ -67,6 +68,7 @@ class TestOnnxConfigRegistration:
             "xlm-roberta",
             "camembert",
             "mpnet",
+            "layoutlm-qa",
             "zoedepth",
         ],
     )
@@ -131,6 +133,59 @@ def test_bert_io_specs_shape_matches(self, bert_config) -> None:
             )
 
 
+class TestLayoutLMQuestionAnsweringOverride:
+    """LayoutLM QA export must include bbox and safe token_type_ids."""
+
+    def test_layoutlm_qa_dummy_inputs_include_bbox_and_zero_token_types(self) -> None:
+        """Dummy inputs must keep bbox while forcing token_type_ids to zero."""
+        from transformers import LayoutLMConfig
+
+        layoutlm_config = LayoutLMConfig(
+            vocab_size=100,
+            hidden_size=64,
+            num_hidden_layers=2,
+            num_attention_heads=2,
+            intermediate_size=128,
+            max_position_embeddings=32,
+            max_2d_position_embeddings=1024,
+            type_vocab_size=1,
+        )
+
+        inputs = generate_dummy_inputs("layoutlm", "question-answering", layoutlm_config)
+
+        assert set(inputs) == {"input_ids", "bbox", "attention_mask", "token_type_ids"}
+        assert inputs["input_ids"].shape == (1, layoutlm_config.max_position_embeddings)
+        assert inputs["bbox"].shape == (1, layoutlm_config.max_position_embeddings, 4)
+        assert inputs["token_type_ids"].shape == (1, layoutlm_config.max_position_embeddings)
+        assert inputs["token_type_ids"].max().item() == 0
+
+    def test_layoutlm_qa_io_specs_include_span_outputs(self) -> None:
+        """LayoutLM QA specs expose document bbox input and span logits outputs."""
+        from transformers import LayoutLMConfig
+
+        layoutlm_config = LayoutLMConfig(
+            vocab_size=100,
+            hidden_size=64,
+            num_hidden_layers=2,
+            num_attention_heads=2,
+            intermediate_size=128,
+            max_position_embeddings=32,
+            max_2d_position_embeddings=1024,
+            type_vocab_size=1,
+        )
+
+        specs = resolve_io_specs("layoutlm", "question-answering", layoutlm_config)
+
+        assert specs["input_names"] == [
+            "input_ids",
+            "bbox",
+            "attention_mask",
+            "token_type_ids",
+        ]
+        assert specs["input_shapes"] == [(1, 32), (1, 32, 4), (1, 32), (1, 32)]
+        assert specs["output_names"] == ["start_logits", "end_logits"]
+
+
 # =============================================================================
 # Class 3: _adjust_position_embeddings unit tests
 # =============================================================================

tests/unit/export/test_pytorch_export.py

@@ -122,6 +122,21 @@ def test_int64_tensor(self) -> None:
         assert t.dtype == torch.int64
         assert (t == 1).all()
 
+    def test_bbox_tensor_generates_ordered_coordinates_within_range(self) -> None:
+        spec = InputTensorSpec(
+            name="bbox",
+            dtype="int32",
+            shape=(2, 128, 4),
+            value_range=(0, 1000),
+        )
+        t = spec.to_tensor()
+        assert t.dtype == torch.int32
+        assert t.shape == (2, 128, 4)
+        assert int(t.min()) >= 0
+        assert int(t.max()) < 1000
+        assert torch.all(t[:, :, 2] >= t[:, :, 0])
+        assert torch.all(t[:, :, 3] >= t[:, :, 1])
+
     def test_no_dtype_defaults_float(self) -> None:
         spec = InputTensorSpec(name="x", shape=(1, 10))
         t = spec.to_tensor()