diff --git a/examples/recipes/DmitrySpartak_layoutlm-invoices/cpu/cpu/document-question-answering_fp16_config.json b/examples/recipes/DmitrySpartak_layoutlm-invoices/cpu/cpu/document-question-answering_fp16_config.json new file mode 100644 index 000000000..44e89502f --- /dev/null +++ b/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" + } +} diff --git a/examples/recipes/DmitrySpartak_layoutlm-invoices/cpu/cpu/document-question-answering_fp32_config.json b/examples/recipes/DmitrySpartak_layoutlm-invoices/cpu/cpu/document-question-answering_fp32_config.json new file mode 100644 index 000000000..99ce8f439 --- /dev/null +++ b/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" + } +} diff --git a/src/winml/modelkit/commands/perf.py b/src/winml/modelkit/commands/perf.py index 30d533699..bfc2d5634 100644 --- a/src/winml/modelkit/commands/perf.py +++ b/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) diff --git a/src/winml/modelkit/loader/resolution.py b/src/winml/modelkit/loader/resolution.py index 548848175..39e0fddb4 100644 --- a/src/winml/modelkit/loader/resolution.py +++ b/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: diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index 3c5b0c12b..454686ecb 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/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", ] diff --git a/src/winml/modelkit/models/hf/layoutlm.py b/src/winml/modelkit/models/hf/layoutlm.py new file mode 100644 index 000000000..0a5a87aab --- /dev/null +++ b/src/winml/modelkit/models/hf/layoutlm.py @@ -0,0 +1,150 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""LayoutLM v1 document-question-answering support. + +LayoutLM v1 uses an extractive question-answering head for document QA: the +model consumes token IDs plus OCR-derived bounding boxes and emits start/end +logits. Optimum does not register this model-type/task pair, and it cannot +infer the task from ``LayoutLMForQuestionAnswering``. + +This module supplies the missing task metadata, model-class mapping, and ONNX +configuration for the whole LayoutLM v1 family. It does not implement OCR, +word/box alignment, or answer decoding; those remain external preprocessing +and postprocessing steps. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from optimum.exporters.onnx.model_configs import LayoutLMOnnxConfig +from optimum.utils.input_generators import DummyTextInputGenerator +from transformers import AutoModelForDocumentQuestionAnswering + +from ...export import register_onnx_overwrite + + +logger = logging.getLogger(__name__) + + +# The task is derived from checkpoint architecture metadata instead of a model +# ID. The resolver consults this registry before Optimum task inference. +ARCHITECTURE_TASK_MAPPING: dict[tuple[str, str], str] = { + ("layoutlm", "LayoutLMForQuestionAnswering"): "document-question-answering", +} + + +# Optimum has no model-class entry for document-question-answering. Transformers +# does: its document-QA auto class maps LayoutLMConfig to +# LayoutLMForQuestionAnswering. +MODEL_CLASS_MAPPING: dict[tuple[str, str], type] = { + ("layoutlm", "document-question-answering"): AutoModelForDocumentQuestionAnswering, +} + + +def _adjust_roberta_position_embeddings(config: Any) -> None: + """Use the non-padding sequence length for RoBERTa-tokenized checkpoints. + + RoBERTa tokenizers reserve ``pad_token_id + 1`` position slots, so their + configs commonly encode ``max_position_embeddings`` as usable length plus + that offset (for example, 514 for a usable length of 512). Only configs + that explicitly declare a RoBERTa tokenizer are adjusted; a LayoutLM model + with a positive padding ID but another position convention is untouched. + """ + if getattr(config, "_position_offset_applied", False): + return + + tokenizer_class = getattr(config, "tokenizer_class", "") or "" + if "roberta" not in tokenizer_class.lower(): + return + + max_positions = getattr(config, "max_position_embeddings", None) + pad_token_id = getattr(config, "pad_token_id", 0) or 0 + if max_positions is None or pad_token_id <= 0: + return + + adjusted = max_positions - pad_token_id - 1 + if adjusted <= 0: + raise ValueError( + "Position offset adjustment would produce non-positive " + f"max_position_embeddings={adjusted} " + f"(original={max_positions}, pad_token_id={pad_token_id})" + ) + + config.max_position_embeddings = adjusted + config._position_offset_applied = True + logger.debug( + "Adjusted LayoutLM max_position_embeddings: %d -> %d (pad_token_id=%d)", + max_positions, + adjusted, + pad_token_id, + ) + + +class LayoutLMTextInputGenerator(DummyTextInputGenerator): # type: ignore[misc] + """Generate token types within the checkpoint's configured vocabulary. + + Optimum's generic text generator samples token-type IDs from ``[0, 2)``. + LayoutLM checkpoints commonly set ``type_vocab_size=1``, where sampling 1 + causes an embedding Gather out of bounds. Derive the exclusive upper bound + from config metadata instead. + """ + + def generate( + self, + input_name: str, + framework: str = "pt", + int_dtype: str = "int64", + float_dtype: str = "fp32", + ) -> Any: + """Generate a safe tensor for ``input_name``.""" + if input_name != "token_type_ids": + return super().generate(input_name, framework, int_dtype, float_dtype) + + type_vocab_size = max(1, int(self.normalized_config.type_vocab_size)) + shape = [self.batch_size, self.sequence_length] + return self.random_int_tensor( + shape, + max_value=type_vocab_size, + min_value=0, + framework=framework, + dtype=int_dtype, + ) + + +@register_onnx_overwrite("layoutlm", "document-question-answering", library_name="transformers") +class LayoutLMDocumentQAOnnxConfig(LayoutLMOnnxConfig): # type: ignore[misc] + """ONNX config for the LayoutLM v1 extractive document-QA head. + + Inputs remain ``input_ids``, ``bbox``, ``attention_mask`` and + ``token_type_ids`` from Optimum's LayoutLM config. Outputs are the two span + tensors returned by ``LayoutLMForQuestionAnswering``. + """ + + DUMMY_INPUT_GENERATOR_CLASSES = ( + LayoutLMTextInputGenerator, + *LayoutLMOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES[1:], + ) + + def __init__(self, config: Any, task: str, **kwargs: Any) -> None: + _adjust_roberta_position_embeddings(config) + super().__init__(config, task, **kwargs) + + @property + def outputs(self) -> dict[str, dict[int, str]]: + """Return the extractive span-head output names and dynamic axes.""" + return { + "start_logits": {0: "batch_size", 1: "sequence_length"}, + "end_logits": {0: "batch_size", 1: "sequence_length"}, + } + + +__all__ = [ + "ARCHITECTURE_TASK_MAPPING", + "MODEL_CLASS_MAPPING", + "LayoutLMDocumentQAOnnxConfig", + "LayoutLMTextInputGenerator", +] diff --git a/tests/unit/commands/test_perf_cli.py b/tests/unit/commands/test_perf_cli.py index be4984977..77d00ec15 100644 --- a/tests/unit/commands/test_perf_cli.py +++ b/tests/unit/commands/test_perf_cli.py @@ -849,6 +849,55 @@ def test_symbolic_override_rejects_non_integer_float_value(self) -> None: ) +class TestGenerateRandomInputs: + @pytest.mark.parametrize("input_name", ["token_type_ids", "arbitrary_integer_input"]) + def test_honors_exclusive_integer_value_range(self, input_name: str) -> None: + """Configured [0, 1) ranges produce only legal zero-valued indices.""" + import numpy as np + + from winml.modelkit.commands.perf import generate_random_inputs + + inputs = generate_random_inputs( + { + "input_names": [input_name], + "input_shapes": [[1, 512]], + "input_types": ["int32"], + "input_value_ranges": {input_name: [0, 1]}, + } + ) + + assert inputs[input_name].shape == (1, 512) + assert inputs[input_name].dtype == np.int64 + assert np.all(inputs[input_name] == 0) + + def test_ranges_and_unspecified_defaults_preserve_existing_generation(self) -> None: + """Ranges are per input; unspecified inputs retain legacy defaults.""" + import numpy as np + + from winml.modelkit.commands.perf import generate_random_inputs + + inputs = generate_random_inputs( + { + "input_names": ["category_ids", "ranged_values", "default_values"], + "input_shapes": [[2, 3], [2, 3], [2, 3]], + "input_types": ["int64", "float16", "float32"], + "input_value_ranges": { + "category_ids": [7, 8], + "ranged_values": [-2.0, -1.0], + }, + } + ) + + assert inputs["category_ids"].shape == (2, 3) + assert np.all(inputs["category_ids"] == 7) + assert inputs["ranged_values"].dtype == np.float32 + assert np.all(inputs["ranged_values"] >= -2.0) + assert np.all(inputs["ranged_values"] < -1.0) + assert inputs["default_values"].dtype == np.float32 + assert np.all(inputs["default_values"] >= 0.0) + assert np.all(inputs["default_values"] < 1.0) + + class TestEffectiveBatchSize: """Throughput must scale by the batch the session actually ran. diff --git a/tests/unit/export/test_layoutlm_onnx_config.py b/tests/unit/export/test_layoutlm_onnx_config.py new file mode 100644 index 000000000..ce6172f5b --- /dev/null +++ b/tests/unit/export/test_layoutlm_onnx_config.py @@ -0,0 +1,91 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for LayoutLM v1 document-question-answering export support.""" + +from __future__ import annotations + +import pytest +from optimum.exporters.tasks import TasksManager +from transformers import LayoutLMConfig + +from winml.modelkit.export import generate_dummy_inputs, resolve_io_specs +from winml.modelkit.models.hf.layoutlm import ( + LayoutLMDocumentQAOnnxConfig, + _adjust_roberta_position_embeddings, +) + + +@pytest.fixture +def layoutlm_config() -> LayoutLMConfig: + config = 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, + ) + config.tokenizer_class = "RobertaTokenizer" + config.architectures = ["LayoutLMForQuestionAnswering"] + return config + + +def test_document_qa_config_registered() -> None: + 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: LayoutLMConfig) -> None: + 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_dummy_token_types_respect_type_vocab_size(layoutlm_config: LayoutLMConfig) -> None: + inputs = generate_dummy_inputs( + "layoutlm", + "document-question-answering", + layoutlm_config, + sequence_length=8, + ) + assert inputs["token_type_ids"].shape == (1, 8) + assert inputs["token_type_ids"].unique().tolist() == [0] + + +def test_roberta_position_offset_is_adjusted_once(layoutlm_config: LayoutLMConfig) -> None: + _adjust_roberta_position_embeddings(layoutlm_config) + assert layoutlm_config.max_position_embeddings == 512 + + _adjust_roberta_position_embeddings(layoutlm_config) + assert layoutlm_config.max_position_embeddings == 512 + + +def test_non_roberta_position_convention_is_untouched() -> None: + config = LayoutLMConfig(max_position_embeddings=512, pad_token_id=1) + config.tokenizer_class = "BertTokenizer" + + _adjust_roberta_position_embeddings(config) + + assert config.max_position_embeddings == 512 + + +def test_roberta_position_offset_rejects_non_positive_length( + layoutlm_config: LayoutLMConfig, +) -> None: + layoutlm_config.max_position_embeddings = 2 + with pytest.raises(ValueError, match="non-positive"): + _adjust_roberta_position_embeddings(layoutlm_config) diff --git a/tests/unit/loader/test_layoutlm_resolution.py b/tests/unit/loader/test_layoutlm_resolution.py new file mode 100644 index 000000000..8411a548e --- /dev/null +++ b/tests/unit/loader/test_layoutlm_resolution.py @@ -0,0 +1,47 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Task and model-class resolution for LayoutLM v1 document QA.""" + +from __future__ import annotations + +from transformers import LayoutLMConfig + +from winml.modelkit.loader.resolution import TaskSource, resolve_task + + +def _document_qa_config() -> LayoutLMConfig: + config = LayoutLMConfig() + config.architectures = ["LayoutLMForQuestionAnswering"] + return config + + +def test_layoutlm_document_qa_resolves_from_architecture() -> None: + resolution = resolve_task(_document_qa_config()) + + assert resolution.task == "document-question-answering" + assert resolution.optimum_task == "document-question-answering" + assert resolution.model_class.__name__ == "AutoModelForDocumentQuestionAnswering" + assert resolution.source == TaskSource.ARCHITECTURE_MAPPING + + +def test_layoutlm_document_qa_explicit_task_uses_document_auto_class() -> None: + resolution = resolve_task( + _document_qa_config(), + task="document-question-answering", + ) + + assert resolution.task == "document-question-answering" + assert resolution.model_class.__name__ == "AutoModelForDocumentQuestionAnswering" + assert resolution.source == TaskSource.USER_TASK + + +def test_layoutlm_other_head_is_not_forced_to_document_qa() -> None: + config = LayoutLMConfig() + config.architectures = ["LayoutLMForTokenClassification"] + + resolution = resolve_task(config) + + assert resolution.task == "token-classification" + assert resolution.source == TaskSource.TASKS_MANAGER