From 44c2a603758226425a96d56153616cc42652d7a7 Mon Sep 17 00:00:00 2001 From: fuliucansheng Date: Thu, 23 Jul 2026 15:49:43 +0800 Subject: [PATCH] fix(layoutlm): resolve question answering heads --- .../cpu/question-answering_fp16_config.json | 104 ++++++++++++++++++ .../cpu/question-answering_fp32_config.json | 81 ++++++++++++++ src/winml/modelkit/loader/resolution.py | 24 +++- tests/unit/loader/test_resolve_task.py | 8 ++ 4 files changed, 216 insertions(+), 1 deletion(-) create mode 100644 examples/recipes/DmitrySpartak_layoutlm-invoices/cpu/cpu/question-answering_fp16_config.json create mode 100644 examples/recipes/DmitrySpartak_layoutlm-invoices/cpu/cpu/question-answering_fp32_config.json diff --git a/examples/recipes/DmitrySpartak_layoutlm-invoices/cpu/cpu/question-answering_fp16_config.json b/examples/recipes/DmitrySpartak_layoutlm-invoices/cpu/cpu/question-answering_fp16_config.json new file mode 100644 index 000000000..47175dac2 --- /dev/null +++ b/examples/recipes/DmitrySpartak_layoutlm-invoices/cpu/cpu/question-answering_fp16_config.json @@ -0,0 +1,104 @@ +{ + "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": { + "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": "question-answering", + "model_id": "DmitrySpartak/layoutlm-invoices", + "model_type": "layoutlm", + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "loader": { + "task": "question-answering", + "model_class": "LayoutLMForQuestionAnswering", + "model_type": "layoutlm" + }, + "compile": null +} diff --git a/examples/recipes/DmitrySpartak_layoutlm-invoices/cpu/cpu/question-answering_fp32_config.json b/examples/recipes/DmitrySpartak_layoutlm-invoices/cpu/cpu/question-answering_fp32_config.json new file mode 100644 index 000000000..0fe7d1a21 --- /dev/null +++ b/examples/recipes/DmitrySpartak_layoutlm-invoices/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/loader/resolution.py b/src/winml/modelkit/loader/resolution.py index 548848175..7bff78859 100644 --- a/src/winml/modelkit/loader/resolution.py +++ b/src/winml/modelkit/loader/resolution.py @@ -178,6 +178,20 @@ def _upgrade_fill_mask_for_seq2seq(task: str, config: PretrainedConfig) -> str: } +_ARCHITECTURE_SUFFIX_TASKS: dict[str, str] = { + "ForQuestionAnswering": "question-answering", +} + + +def _infer_task_from_architecture_suffix(config: PretrainedConfig) -> str | None: + """Infer task from architecture head suffix when Optimum has no mapping.""" + for arch_name in getattr(config, "architectures", None) or []: + for suffix, task in _ARCHITECTURE_SUFFIX_TASKS.items(): + if arch_name.endswith(suffix): + return task + return None + + def _resolve_task_modality(config: PretrainedConfig, task: str) -> str: """Upgrade a modality-blind ``feature-extraction`` to its modality-aware variant. @@ -415,8 +429,16 @@ def _infer_task_from_architecture(config: PretrainedConfig) -> str: Includes the encoder-decoder fill-mask -> text2text-generation correction. """ + model_class = _resolve_model_class_from_config(config) + try: + task = _detect_task_from_model_class(model_class) + except ValueError: + suffix_task = _infer_task_from_architecture_suffix(config) + if suffix_task is None: + raise + task = suffix_task return _upgrade_fill_mask_for_seq2seq( - _detect_task_from_model_class(_resolve_model_class_from_config(config)), + task, config, ) diff --git a/tests/unit/loader/test_resolve_task.py b/tests/unit/loader/test_resolve_task.py index fffd20fe1..b9999571a 100644 --- a/tests/unit/loader/test_resolve_task.py +++ b/tests/unit/loader/test_resolve_task.py @@ -48,6 +48,14 @@ def test_seq2seq_fill_mask_is_upgraded_and_source_is_tasks_manager(): assert r.composite == {"encoder": "feature-extraction", "decoder": "text2text-generation"} +def test_architecture_suffix_fallback_resolves_layoutlm_question_answering(): + r = resolve_task(_cfg("layoutlm", ["LayoutLMForQuestionAnswering"])) + assert r.task == "question-answering" + assert r.optimum_task == "question-answering" + assert r.model_class.__name__ == "AutoModelForQuestionAnswering" + assert r.source == TaskSource.TASKS_MANAGER + + def test_user_task_preserved_verbatim_no_modality_upgrade(): r = resolve_task(_cfg("vit", ["ViTModel"]), task="feature-extraction") assert r.task == "feature-extraction"