diff --git a/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp16_config.json b/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp16_config.json new file mode 100644 index 000000000..0fe7d1a21 --- /dev/null +++ b/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" + } +} diff --git a/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp32_config.json b/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp32_config.json new file mode 100644 index 000000000..0fe7d1a21 --- /dev/null +++ b/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" + } +} diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index 094714494..13250616f 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/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 diff --git a/src/winml/modelkit/models/hf/layoutlm.py b/src/winml/modelkit/models/hf/layoutlm.py new file mode 100644 index 000000000..8085eb0ef --- /dev/null +++ b/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, + ) diff --git a/src/winml/modelkit/onnx/io.py b/src/winml/modelkit/onnx/io.py index 1c58362cf..d2ad85878 100644 --- a/src/winml/modelkit/onnx/io.py +++ b/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) diff --git a/tests/unit/export/test_onnx_config_overrides.py b/tests/unit/export/test_onnx_config_overrides.py index f7ba15a39..29478512c 100644 --- a/tests/unit/export/test_onnx_config_overrides.py +++ b/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 # ============================================================================= diff --git a/tests/unit/export/test_pytorch_export.py b/tests/unit/export/test_pytorch_export.py index 5b6aeae26..f22fc47b6 100644 --- a/tests/unit/export/test_pytorch_export.py +++ b/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()