Add LayoutLM v1 document-question-answering support (impira/layoutlm-invoices)

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

@@ -0,0 +1,99 @@
+{
+  "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,
+          1
+        ]
+      },
+      {
+        "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": {},
+  "quant": {
+    "mode": "fp16",
+    "samples": 10,
+    "calibration_method": "minmax",
+    "weight_type": "uint8",
+    "activation_type": "uint8",
+    "per_channel": false,
+    "symmetric": false,
+    "weight_symmetric": null,
+    "activation_symmetric": null,
+    "save_calibration": false,
+    "distribution": "uniform",
+    "seed": null,
+    "calibration_load_path": null,
+    "calibration_save_path": null,
+    "op_types_to_quantize": null,
+    "nodes_to_exclude": null,
+    "fp16_keep_io_types": true,
+    "fp16_op_block_list": null
+  },
+  "compile": null,
+  "loader": {
+    "task": "document-question-answering",
+    "model_class": "AutoModelForDocumentQuestionAnswering",
+    "model_type": "layoutlm"
+  }
+}

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

@@ -0,0 +1,82 @@
+{
+    "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,
+                                                                 1
+                                                             ]
+                                         },
+                                         {
+                                             "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":  {
+
+              },
+    "quant":  null,
+    "compile":  null,
+    "loader":  {
+                   "task":  "document-question-answering",
+                   "model_class":  "AutoModelForDocumentQuestionAnswering",
+                   "model_type":  "layoutlm"
+               }
+}

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

@@ -48,6 +48,9 @@
 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 (
+    LayoutLMDocumentQAOnnxConfig as _LayoutLMDocumentQAOnnxConfig,  # 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,111 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+"""LayoutLM (v1) document-question-answering ONNX export config.
+
+Optimum's TasksManager registers ``layoutlm`` only for feature-extraction,
+fill-mask, text-classification and token-classification -- it has no
+document-question-answering config. However, transformers'
+``LayoutLMForQuestionAnswering`` (the head behind
+``AutoModelForDocumentQuestionAnswering`` for layoutlm v1, e.g.
+``impira/layoutlm-invoices``) is an *extractive* QA head that returns
+``start_logits`` / ``end_logits``.
+
+This module registers a ``document-question-answering`` OnnxConfig for
+layoutlm that reuses ``LayoutLMOnnxConfig``'s inputs (``input_ids``, ``bbox``,
+``attention_mask``, ``token_type_ids``) and overrides the outputs to the
+extractive-QA pair. Optimum's default document-question-answering output is a
+single ``logits`` tensor, which does not match this head; the override emits
+the correct ``start_logits`` / ``end_logits`` pair.
+
+This module provides:
+- LayoutLMDocumentQAOnnxConfig: ONNX export config for layoutlm document-QA
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from optimum.exporters.onnx.model_configs import LayoutLMOnnxConfig
+
+from ...export import register_onnx_overwrite
+
+
+logger = logging.getLogger(__name__)
+
+
+def _adjust_position_embeddings(config: Any) -> None:
+    """Adjust max_position_embeddings for RoBERTa-style position offset.
+
+    Some LayoutLM checkpoints (e.g. ``impira/layoutlm-invoices``) are built on a
+    RoBERTa vocabulary/tokenizer and inherit its position-embedding convention::
+
+        max_position_embeddings = usable_length + pad_token_id + 1
+
+    e.g. 514 = 512 + 1 + 1 (pad_token_id=1). LayoutLM's ``sequence_length`` is
+    derived from ``max_position_embeddings``; using the raw value causes a
+    position index out-of-bounds ("index out of range in self") during ONNX
+    export tracing. This adjusts ``config.max_position_embeddings`` in-place to
+    the usable length. A sentinel prevents double-adjustment on config reuse.
+    """
+    if getattr(config, "_position_offset_applied", False):
+        return
+
+    if not hasattr(config, "max_position_embeddings"):
+        return
+
+    pad_token_id = getattr(config, "pad_token_id", 0) or 0
+    if pad_token_id > 0:
+        original = config.max_position_embeddings
+        adjusted = original - pad_token_id - 1
+        if adjusted <= 0:
+            raise ValueError(
+                f"Position offset adjustment would produce non-positive "
+                f"max_position_embeddings={adjusted} "
+                f"(original={original}, pad_token_id={pad_token_id})"
+            )
+        config.max_position_embeddings = adjusted
+        config._position_offset_applied = True
+        logger.debug(
+            "Adjusted max_position_embeddings: %d -> %d (pad_token_id=%d)",
+            original,
+            adjusted,
+            pad_token_id,
+        )
+
+
+@register_onnx_overwrite("layoutlm", "document-question-answering", library_name="transformers")
+class LayoutLMDocumentQAOnnxConfig(LayoutLMOnnxConfig):  # type: ignore[misc]  # optimum base is untyped
+    """LayoutLM document-question-answering (extractive) OnnxConfig.
+
+    Inputs (inherited from LayoutLMOnnxConfig):
+        - input_ids: {0: "batch_size", 1: "sequence_length"}
+        - bbox: {0: "batch_size", 1: "sequence_length"}
+        - attention_mask: {0: "batch_size", 1: "sequence_length"}
+        - token_type_ids: {0: "batch_size", 1: "sequence_length"}
+
+    Outputs (extractive QA -- overrides Optimum's single-``logits`` default):
+        - start_logits: {0: "batch_size", 1: "sequence_length"}
+        - end_logits: {0: "batch_size", 1: "sequence_length"}
+
+    RoBERTa-vocab LayoutLM checkpoints set ``max_position_embeddings`` to the
+    usable length plus a pad-token offset; the constructor adjusts it back to
+    the usable length so ONNX export tracing does not hit a position-index
+    out-of-bounds. LayoutLM's own dummy-input generators (which supply ``bbox``)
+    are preserved.
+    """
+
+    def __init__(self, config: Any, task: str, **kwargs: Any) -> None:
+        _adjust_position_embeddings(config)
+        super().__init__(config, task, **kwargs)
+
+    @property
+    def outputs(self) -> dict[str, dict[int, str]]:
+        """Extractive-QA outputs: ``start_logits`` and ``end_logits``."""
+        return {
+            "start_logits": {0: "batch_size", 1: "sequence_length"},
+            "end_logits": {0: "batch_size", 1: "sequence_length"},
+        }

tests/unit/export/test_layoutlm_onnx_config.py

@@ -0,0 +1,78 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Tests for LayoutLM v1 document-question-answering export registration."""
+
+from __future__ import annotations
+
+import pytest
+from optimum.exporters.tasks import TasksManager
+
+from winml.modelkit.export import resolve_io_specs
+from winml.modelkit.models.hf.layoutlm import (
+    LayoutLMDocumentQAOnnxConfig,
+    _adjust_position_embeddings,
+)
+
+
+@pytest.fixture
+def layoutlm_config():
+    """Return a small LayoutLM config with RoBERTa-style position offset."""
+    from transformers import LayoutLMConfig
+
+    return LayoutLMConfig(
+        vocab_size=100,
+        hidden_size=32,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        intermediate_size=64,
+        max_position_embeddings=514,
+        pad_token_id=1,
+        type_vocab_size=1,
+    )
+
+
+def test_document_qa_config_registered() -> None:
+    """The document-QA task resolves to the WinML LayoutLM override."""
+    constructor = TasksManager.get_exporter_config_constructor(
+        exporter="onnx",
+        model_type="layoutlm",
+        task="document-question-answering",
+        library_name="transformers",
+    )
+    assert constructor.func is LayoutLMDocumentQAOnnxConfig
+
+
+def test_document_qa_io_specs(layoutlm_config) -> None:
+    """LayoutLM document-QA keeps all four inputs and emits two logits."""
+    specs = resolve_io_specs(
+        "layoutlm",
+        "document-question-answering",
+        layoutlm_config,
+    )
+    assert specs["input_names"] == [
+        "input_ids",
+        "bbox",
+        "attention_mask",
+        "token_type_ids",
+    ]
+    assert specs["output_names"] == ["start_logits", "end_logits"]
+
+
+def test_position_offset_is_adjusted_once(layoutlm_config) -> None:
+    """RoBERTa-style 514 positions become a usable sequence length of 512."""
+    _adjust_position_embeddings(layoutlm_config)
+    assert layoutlm_config.max_position_embeddings == 512
+
+    _adjust_position_embeddings(layoutlm_config)
+    assert layoutlm_config.max_position_embeddings == 512
+
+
+def test_position_offset_rejects_non_positive_length(layoutlm_config) -> None:
+    """Invalid position metadata fails explicitly instead of tracing bad indices."""
+    layoutlm_config.max_position_embeddings = 2
+    layoutlm_config.pad_token_id = 1
+
+    with pytest.raises(ValueError, match="non-positive"):
+        _adjust_position_embeddings(layoutlm_config)