feat(layoutlm): add document question answering support

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

@@ -0,0 +1,102 @@
+{
+  "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,
+    "task": "document-question-answering",
+    "model_id": "DmitrySpartak/layoutlm-invoices",
+    "model_type": "layoutlm",
+    "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/DmitrySpartak_layoutlm-invoices/cpu/cpu/document-question-answering_fp32_config.json

@@ -0,0 +1,80 @@
+{
+  "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/commands/perf.py

@@ -262,6 +262,7 @@ def generate_random_inputs(
     symbolic_shapes = io_config.get("input_symbolic_shapes") or [
         [None] * len(s) for s in io_config["input_shapes"]
     ]
+    value_ranges = io_config.get("input_value_ranges") or {}
     overrides = shape_config or {}
 
     specs: dict[str, dict[str, Any]] = {}
@@ -290,6 +291,11 @@ def generate_random_inputs(
             "dtype": gen_dtype,
             "shape": list(resolved_shape),
         }
+        if name in value_ranges:
+            low, high = value_ranges[name]
+            # InputTensorSpec/config ranges are [low, high), while the legacy
+            # manual NumPy generator treats an integer spec's maximum as inclusive.
+            specs[name]["range"] = [int(low), int(high) - 1] if gen_dtype == "int" else [low, high]
 
     return generate_dummy_inputs_from_specs(specs)
 

src/winml/modelkit/loader/resolution.py

@@ -251,6 +251,7 @@ class TaskSource(str, Enum):
     USER_CLASS = "user-class"  # user passed --model-class; task inferred
     MODEL_ID_DEFAULT = "model-id-default"  # MODEL_TASK_MAPPING model-id default
     SENTINEL_DEFAULT = "sentinel-default"  # (model_type, None) sentinel
+    ARCHITECTURE_MAPPING = "architecture-mapping"  # models/hf metadata override
     TASKS_MANAGER = "tasks-manager"  # Optimum inference (incl. fill-mask upgrade)
     WRAPPED_LIBRARY = "wrapped-library"  # no architectures -> first supported task
     PIPELINE_TAG = "pipeline-tag"  # Hub pipeline_tag fallback
@@ -421,6 +422,25 @@ def _infer_task_from_architecture(config: PretrainedConfig) -> str:
     )
 
 
+def _get_architecture_task_override(config: PretrainedConfig) -> str | None:
+    """Return a registered task derived from model type + architecture head.
+
+    This is the data-driven escape hatch for concrete HuggingFace heads that
+    Optimum cannot infer (or infers incorrectly). Registrations live beside the
+    architecture's ONNX config under ``models/hf/``; no checkpoint IDs or
+    architecture branches belong in this shared resolver.
+    """
+    architectures = getattr(config, "architectures", None)
+    model_type = getattr(config, "model_type", None)
+    if not architectures or not model_type:
+        return None
+
+    from ..models.hf import ARCHITECTURE_TASK_MAPPING
+
+    model_type_norm = model_type.lower().replace("_", "-")
+    return ARCHITECTURE_TASK_MAPPING.get((model_type_norm, architectures[0]))
+
+
 def _composite_components_for_task(model_type: str, task: str) -> CompositeComponents | None:
     """Composite components serving a *detected* task, else None.
 
@@ -513,7 +533,9 @@ def resolve_task(
             # Task inferred from the architecture: surface it modality-aware, consistent
             # with the detection path (Stage 3), so e.g. a ViT backbone is
             # image-feature-extraction rather than the modality-blind feature-extraction.
-            opt_task = _infer_task_from_architecture(config)
+            opt_task = _get_architecture_task_override(config) or _infer_task_from_architecture(
+                config
+            )
             surfaced = _resolve_task_modality(config, opt_task)
         # A WinML build variant (model_type_override) may name a custom wrapper
         # registered in MODEL_CLASS_MAPPING rather than a transformers class —
@@ -604,13 +626,18 @@ def resolve_task(
             # lookup failure here — e.g. a wrapped library whose classes aren't registered
             # under framework="pt" — can't escape as a raw KeyError.
 
-    # 1c. TasksManager (reads config.architectures)
+    # 1c. Architecture metadata override, then TasksManager (both read
+    #     config.architectures). The override handles heads Optimum does not map.
     if opt_task is None:
-        try:
-            opt_task = _infer_task_from_architecture(config)
-            source = TaskSource.TASKS_MANAGER
-        except ValueError:
-            opt_task = None
+        opt_task = _get_architecture_task_override(config)
+        if opt_task is not None:
+            source = TaskSource.ARCHITECTURE_MAPPING
+        else:
+            try:
+                opt_task = _infer_task_from_architecture(config)
+                source = TaskSource.TASKS_MANAGER
+            except ValueError:
+                opt_task = None
 
     # 1d. Hub pipeline_tag fallback
     if opt_task is None and model_id and model_type:

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

@@ -48,6 +48,11 @@
 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 ARCHITECTURE_TASK_MAPPING as _LAYOUTLM_ARCHITECTURE_TASK_MAPPING
+from .layoutlm import MODEL_CLASS_MAPPING as _LAYOUTLM_CLASS_MAPPING
+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
@@ -120,6 +125,7 @@
         _BART_CLASS_MAPPING,
         _BLIP_CLASS_MAPPING,
         _CLIP_CLASS_MAPPING,
+        _LAYOUTLM_CLASS_MAPPING,
         _MARIAN_CLASS_MAPPING,
         _MU2_CLASS_MAPPING,
         _QWEN_CLASS_MAPPING,
@@ -136,6 +142,13 @@
     for _key, _model_cls in _sub_mapping.items()
 }
 
+# Registry for architecture heads that Optimum cannot map to their pipeline
+# task. Keys include model_type so identically named custom classes cannot
+# collide across model families.
+ARCHITECTURE_TASK_MAPPING: dict[tuple[str, str], str] = {
+    **_LAYOUTLM_ARCHITECTURE_TASK_MAPPING,
+}
+
 # Registry: model_type -> WinMLBuildConfig
 # Only models that need non-autoconf-discoverable settings retain configs.
 # Models with only optim flags rely on the analyzer autoconf loop.
@@ -164,6 +177,7 @@
 }
 
 __all__ = [
+    "ARCHITECTURE_TASK_MAPPING",
     "MODEL_BUILD_CONFIGS",
     "MODEL_CLASS_MAPPING",
 ]