diff --git a/examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp16_config.json b/examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp16_config.json new file mode 100644 index 000000000..5dcd84f3e --- /dev/null +++ b/examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp16_config.json @@ -0,0 +1,53 @@ +{ + "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": "pixel_values", + "dtype": "float32", + "shape": [1, 3, 32, 128], + "value_range": [0, 1] + } + ], + "output_tensors": [ + {"name": "char_logits"}, + {"name": "bpe_logits"}, + {"name": "wp_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": "image-to-text", + "model_class": "MgpstrForSceneTextRecognition", + "model_type": "mgp-str" + } +} diff --git a/examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp32_config.json b/examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp32_config.json new file mode 100644 index 000000000..38b8fc6e3 --- /dev/null +++ b/examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp32_config.json @@ -0,0 +1,34 @@ +{ + "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": "pixel_values", + "dtype": "float32", + "shape": [1, 3, 32, 128], + "value_range": [0, 1] + } + ], + "output_tensors": [ + {"name": "char_logits"}, + {"name": "bpe_logits"}, + {"name": "wp_logits"} + ] + }, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "image-to-text", + "model_class": "MgpstrForSceneTextRecognition", + "model_type": "mgp-str" + } +} diff --git a/src/winml/modelkit/eval/image_to_text_evaluator.py b/src/winml/modelkit/eval/image_to_text_evaluator.py index 9f541f173..97bbd2360 100644 --- a/src/winml/modelkit/eval/image_to_text_evaluator.py +++ b/src/winml/modelkit/eval/image_to_text_evaluator.py @@ -29,6 +29,7 @@ if TYPE_CHECKING: from datasets import Dataset + from transformers.pipelines.base import Pipeline from ..models.winml.base import WinMLPreTrainedModel from .config import DatasetConfig, WinMLEvaluationConfig @@ -56,6 +57,17 @@ def align_labels(self, dataset: Dataset, ds_config: DatasetConfig) -> Dataset: """No-op: free-text labels need no ClassLabel alignment.""" return dataset + def prepare_pipeline(self) -> Pipeline: + """Create an image-to-text pipeline, including registered specializations.""" + from typing import cast + + from ..inference.pipeline import create_pipeline + + return cast( + "Pipeline", + create_pipeline("image-to-text", self.model, self.config.model_id), + ) + def compute(self) -> dict[str, Any]: """Run the pipeline over each sample and return CER + CIDEr.""" from tqdm.auto import tqdm diff --git a/src/winml/modelkit/inference/pipeline.py b/src/winml/modelkit/inference/pipeline.py index fc8bdb66c..f3ff56ae1 100644 --- a/src/winml/modelkit/inference/pipeline.py +++ b/src/winml/modelkit/inference/pipeline.py @@ -20,7 +20,7 @@ import inspect import logging -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast if TYPE_CHECKING: @@ -37,6 +37,10 @@ "sentence-similarity": "feature-extraction", } +_SPECIALIZED_PIPELINE_REGISTRY: dict[tuple[str, str], str] = { + ("mgp-str", "image-to-text"): "MgpstrImageToTextPipeline", +} + def create_pipeline( task: str, @@ -57,6 +61,23 @@ def create_pipeline( Returns: A configured ``transformers.Pipeline`` ready for inference. """ + model_type = getattr(getattr(model, "config", None), "model_type", None) + if not isinstance(model_type, str): + model_type = "" + specialized = _SPECIALIZED_PIPELINE_REGISTRY.get((model_type, task)) + if specialized is not None: + if model_id is None: + raise ValueError(f"{specialized} requires a model ID to load its processor.") + from ..models.winml.image_to_text import MgpstrImageToTextPipeline + + specialized_classes = {"MgpstrImageToTextPipeline": MgpstrImageToTextPipeline} + pipe = specialized_classes[specialized]( + cast("WinMLPreTrainedModel", model), model_id + ) + _adapt_image_processor_size(pipe, task, model) + logger.info("Created specialized pipeline: task=%s model=%s", task, model_id) + return pipe + from transformers import pipeline kwargs: dict[str, Any] = { diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index 3c5b0c12b..80c80229d 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/src/winml/modelkit/models/hf/__init__.py @@ -52,6 +52,10 @@ from .marian import MODEL_CLASS_MAPPING as _MARIAN_CLASS_MAPPING from .marian import MarianDecoderIOConfig as _MarianDecoderIOConfig # triggers registration from .marian import MarianEncoderIOConfig as _MarianEncoderIOConfig # triggers registration +from .mgp_str import MODEL_CLASS_MAPPING as _MGPSTR_CLASS_MAPPING +from .mgp_str import ( + MgpstrImage2TextOnnxConfig as _MgpstrImage2TextOnnxConfig, # triggers registration +) from .mu2 import MODEL_CLASS_MAPPING as _MU2_CLASS_MAPPING from .mu2 import MU2_CONFIG from .mu2 import Mu2DecoderIOConfig as _Mu2DecoderIOConfig # triggers registration @@ -121,6 +125,7 @@ _BLIP_CLASS_MAPPING, _CLIP_CLASS_MAPPING, _MARIAN_CLASS_MAPPING, + _MGPSTR_CLASS_MAPPING, _MU2_CLASS_MAPPING, _QWEN_CLASS_MAPPING, _QWEN_TO_CLASS_MAPPING, diff --git a/src/winml/modelkit/models/hf/mgp_str.py b/src/winml/modelkit/models/hf/mgp_str.py new file mode 100644 index 000000000..3118c329e --- /dev/null +++ b/src/winml/modelkit/models/hf/mgp_str.py @@ -0,0 +1,59 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Configure MGP-STR for Hugging Face scene-text-recognition export. + +MGP-STR is a Vision Transformer-based scene text recognition (STR) model. The +upstream ``MgpstrForSceneTextRecognition`` head produces three logit tensors — +``char_logits``, ``bpe_logits``, ``wp_logits`` — at three granularities +(character / byte-pair / word-piece), which the ``MgpstrProcessor`` combines +into the final decoded string. + +The vendor ``MgpstrOnnxConfig`` (Optimum) already exposes the 3-head outputs +correctly but is registered ONLY under the ``feature-extraction`` task. End +users naturally reach for the ``image-to-text`` task label for STR work; this +module registers the same export config under ``image-to-text`` so the +user-facing task resolves cleanly. + +Registering the canonical task lets the shared resolver accept the published +``image-to-text`` pipeline metadata even when checkpoint architecture metadata +is stale, then bind the family to its concrete head-bearing class. +""" + +from __future__ import annotations + +from optimum.exporters.onnx.model_configs import MgpstrOnnxConfig +from transformers import MgpstrForSceneTextRecognition + +from ...export import register_onnx_overwrite + + +# ============================================================================= +# Image-to-text alias for MGP-STR +# ============================================================================= + + +@register_onnx_overwrite("mgp-str", "image-to-text", library_name="transformers") +class MgpstrImage2TextOnnxConfig(MgpstrOnnxConfig): # type: ignore[misc] + """MGP-STR ONNX config bound to the ``image-to-text`` task. + + The 3-head ``(char_logits, bpe_logits, wp_logits)`` output contract and + the ``pixel_values`` input contract are inherited unchanged from + ``MgpstrOnnxConfig``. The only purpose of this subclass is to register + the same export semantics under the ``image-to-text`` task name so users + can build MGP-STR with the natural task label. + """ + + +# ============================================================================= +# Model Class Mapping +# ============================================================================= + +# (model_type, task) -> HF model class. Binds the ``image-to-text`` task on +# MGP-STR to ``MgpstrForSceneTextRecognition`` (the head-bearing class with the +# 3-granularity outputs), instead of letting the loader fall back to +# ``AutoModelForVision2Seq`` — MGP-STR is NOT a Vision2Seq architecture. +MODEL_CLASS_MAPPING: dict[tuple[str, str], type] = { + ("mgp-str", "image-to-text"): MgpstrForSceneTextRecognition, +} diff --git a/src/winml/modelkit/models/winml/__init__.py b/src/winml/modelkit/models/winml/__init__.py index bd1ac303b..4c073eceb 100644 --- a/src/winml/modelkit/models/winml/__init__.py +++ b/src/winml/modelkit/models/winml/__init__.py @@ -59,6 +59,7 @@ # - Specific OPSET requirements # - Input remapping (non-standard tensor names) # - Custom pre/post-processing + ("mgp-str", "image-to-text"): "WinMLModelForMgpstrSceneTextRecognition", } @@ -82,6 +83,7 @@ def _import_winml_class(class_name: str) -> type[WinMLPreTrainedModel]: WinMLModelForImageSegmentation, WinMLModelForSemanticSegmentation, ) + from .image_to_text import WinMLModelForMgpstrSceneTextRecognition from .object_detection import WinMLModelForObjectDetection from .question_answering import WinMLModelForQuestionAnswering from .sequence_classification import WinMLModelForSequenceClassification @@ -91,6 +93,7 @@ def _import_winml_class(class_name: str) -> type[WinMLPreTrainedModel]: "WinMLModelForDepthEstimation": WinMLModelForDepthEstimation, "WinMLModelForFeatureExtraction": WinMLModelForFeatureExtraction, "WinMLModelForImageClassification": WinMLModelForImageClassification, + "WinMLModelForMgpstrSceneTextRecognition": WinMLModelForMgpstrSceneTextRecognition, "WinMLModelForImageSegmentation": WinMLModelForImageSegmentation, "WinMLModelForObjectDetection": WinMLModelForObjectDetection, "WinMLModelForQuestionAnswering": WinMLModelForQuestionAnswering, @@ -204,6 +207,7 @@ def register_specialization(model_type: str, task: str, class_name: str) -> None WinMLModelForImageSegmentation, WinMLModelForSemanticSegmentation, ) +from .image_to_text import WinMLModelForMgpstrSceneTextRecognition from .kv_cache import ( WinMLCache, WinMLSlidingWindowCache, @@ -233,6 +237,7 @@ def register_specialization(model_type: str, task: str, class_name: str) -> None "WinMLModelForGenericTask", "WinMLModelForImageClassification", "WinMLModelForImageSegmentation", + "WinMLModelForMgpstrSceneTextRecognition", "WinMLModelForObjectDetection", "WinMLModelForSemanticSegmentation", "WinMLModelForSequenceClassification", diff --git a/src/winml/modelkit/models/winml/image_to_text.py b/src/winml/modelkit/models/winml/image_to_text.py new file mode 100644 index 000000000..85e41e6bd --- /dev/null +++ b/src/winml/modelkit/models/winml/image_to_text.py @@ -0,0 +1,76 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Image-to-text inference wrappers and pipelines.""" + +from __future__ import annotations + +from typing import Any, cast + +from .base import WinMLPreTrainedModel + + +class WinMLModelForMgpstrSceneTextRecognition(WinMLPreTrainedModel): + """Expose MGP-STR's three ONNX heads in the Transformers output contract.""" + + main_input_name = "pixel_values" + + def forward(self, **kwargs: Any) -> Any: + """Run the ONNX graph and return its ordered three-head logits tuple.""" + from transformers.models.mgp_str.modeling_mgp_str import MgpstrModelOutput + + try: + pixel_values = kwargs.pop("pixel_values") + except KeyError as exc: + raise ValueError("MGP-STR inference requires 'pixel_values'.") from exc + + outputs = self._run_inference(self._format_inputs(pixel_values=pixel_values)) + return MgpstrModelOutput( + # Transformers types this heterogeneous three-head tuple as a + # one-item FloatTensor tuple even though that contradicts its + # runtime MGP-STR contract. + logits=cast( + "Any", + ( + outputs["char_logits"], + outputs["bpe_logits"], + outputs["wp_logits"], + ), + ) + ) + + +class MgpstrImageToTextPipeline: + """Preprocess images and decode MGP-STR's character/BPE/WordPiece heads.""" + + def __init__(self, model: WinMLPreTrainedModel, model_id: str) -> None: + from transformers import AutoProcessor + + self.model = model + self.processor = AutoProcessor.from_pretrained(model_id) + self.image_processor = self.processor.image_processor + self.tokenizer = getattr(self.processor, "char_tokenizer", None) + self._preprocess_params: dict[str, Any] = {} + + def _sanitize_parameters(self, **_kwargs: Any) -> tuple[dict, dict, dict]: + """Match the pipeline introspection contract used by inference clients.""" + return {}, {}, {} + + def __call__(self, images: Any, *, prompt: str | None = None, **_kwargs: Any) -> Any: + """Recognize text in one image or a batch of images.""" + if prompt is not None: + raise ValueError("MGP-STR scene text recognition does not accept a text prompt.") + + model_inputs = self.processor(images=images, return_tensors="pt") + outputs = self.model(**model_inputs) + decoded = self.processor.batch_decode(outputs.logits) + + records = [] + for index, text in enumerate(decoded["generated_text"]): + record: dict[str, Any] = {"generated_text": text} + if "scores" in decoded: + score = decoded["scores"][index] + record["score"] = float(score.item() if hasattr(score, "item") else score) + records.append(record) + return records diff --git a/tests/unit/eval/test_image_to_text_evaluator.py b/tests/unit/eval/test_image_to_text_evaluator.py index f20a3665b..5bb593979 100644 --- a/tests/unit/eval/test_image_to_text_evaluator.py +++ b/tests/unit/eval/test_image_to_text_evaluator.py @@ -78,6 +78,29 @@ def test_registered(self): ) +class TestPreparePipeline: + def test_uses_shared_pipeline_factory(self): + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + evaluator = object.__new__(WinMLImageToTextEvaluator) + evaluator.model = MagicMock() + evaluator.config = WinMLEvaluationConfig( + model_id="local/mgp-str", + task="image-to-text", + dataset=DatasetConfig(path="local/dataset"), + ) + expected = MagicMock() + + with patch( + "winml.modelkit.inference.pipeline.create_pipeline", + return_value=expected, + ) as create: + result = evaluator.prepare_pipeline() + + assert result is expected + create.assert_called_once_with("image-to-text", evaluator.model, "local/mgp-str") + + class TestCompute: """compute() iterates samples through the pipeline and aggregates metrics.""" diff --git a/tests/unit/export/test_mgp_str_onnx_config.py b/tests/unit/export/test_mgp_str_onnx_config.py new file mode 100644 index 000000000..71b2783a5 --- /dev/null +++ b/tests/unit/export/test_mgp_str_onnx_config.py @@ -0,0 +1,128 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Tests for MGP-STR image-to-text ONNX config registration and I/O specs. + +MGP-STR is a scene-text-recognition model whose vendor ``MgpstrOnnxConfig`` +(Optimum) is registered only under ``feature-extraction``. This module +verifies the contribution's task-label alias — ``MgpstrImage2TextOnnxConfig`` +registered under ``image-to-text`` — and locks in the inherited I/O contract: + +- Input: pixel_values +- Outputs: char_logits, bpe_logits, wp_logits (3 granularity heads) + +It also verifies the MODEL_CLASS_MAPPING binding to +``MgpstrForSceneTextRecognition`` (MGP-STR is NOT a Vision2Seq model, so the +loader must not fall back to AutoModelForVision2Seq). + +See also: modelkit/models/hf/mgp_str.py +""" + +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.mgp_str import ( + MODEL_CLASS_MAPPING, + MgpstrImage2TextOnnxConfig, +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def mgp_str_config(): + """Return a small MGP-STR config with the production I/O topology.""" + from transformers import MgpstrConfig + + return MgpstrConfig( + image_size=[32, 128], + patch_size=4, + num_channels=3, + hidden_size=48, + num_hidden_layers=1, + num_attention_heads=2, + mlp_ratio=2, + max_token_length=27, + num_character_labels=38, + num_bpe_labels=50257, + num_wordpiece_labels=30522, + ) + + +# ============================================================================= +# Registration verification +# ============================================================================= + + +class TestMgpStrRegistration: + """Verify MgpstrImage2TextOnnxConfig registration and inheritance.""" + + def test_mgp_str_config_registered(self) -> None: + """MgpstrImage2TextOnnxConfig must be registered for image-to-text.""" + config_constructor = TasksManager.get_exporter_config_constructor( + exporter="onnx", + model_type="mgp-str", + task="image-to-text", + library_name="transformers", + ) + actual_class_name = config_constructor.func.__name__ + assert actual_class_name == "MgpstrImage2TextOnnxConfig", ( + f"Expected MgpstrImage2TextOnnxConfig for mgp-str/image-to-text, " + f"got {actual_class_name}. Registration may have failed." + ) + + def test_vendor_model_patcher_is_preserved(self) -> None: + """Subclassing the vendor config must retain its export-time patcher.""" + from optimum.exporters.onnx.model_configs import MgpstrOnnxConfig + + assert MgpstrImage2TextOnnxConfig._MODEL_PATCHER is MgpstrOnnxConfig._MODEL_PATCHER + + +# ============================================================================= +# I/O verification +# ============================================================================= + + +class TestMgpStrIOSpecs: + """Verify I/O specs for MGP-STR scene text recognition.""" + + def test_input_is_pixel_values(self, mgp_str_config) -> None: + """Input must be the single pixel_values tensor.""" + specs = resolve_io_specs("mgp-str", "image-to-text", mgp_str_config) + assert specs["input_names"] == ["pixel_values"] + + def test_outputs_are_three_heads(self, mgp_str_config) -> None: + """Outputs must be the three granularity logit heads.""" + specs = resolve_io_specs("mgp-str", "image-to-text", mgp_str_config) + assert set(specs["output_names"]) == { + "char_logits", + "bpe_logits", + "wp_logits", + } + + +# ============================================================================= +# Model class mapping verification +# ============================================================================= + + +class TestMgpStrModelClassMapping: + """Verify the model-type/task tuple binds to the head-bearing HF class.""" + + def test_mapping_binds_scene_text_head(self) -> None: + """MGP-STR image-to-text must map to its recognition head.""" + from transformers import MgpstrForSceneTextRecognition + + assert ( + MODEL_CLASS_MAPPING[("mgp-str", "image-to-text")] + is MgpstrForSceneTextRecognition + ) + assert set(MODEL_CLASS_MAPPING) == {("mgp-str", "image-to-text")} diff --git a/tests/unit/inference/test_pipeline.py b/tests/unit/inference/test_pipeline.py index 154e96887..c93936f89 100644 --- a/tests/unit/inference/test_pipeline.py +++ b/tests/unit/inference/test_pipeline.py @@ -15,13 +15,14 @@ import inspect from typing import Any -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch from winml.modelkit.inference.pipeline import ( _HF_PIPELINE_TASK_MAP, _adapt_image_processor_size, _adapt_tokenizer_padding, _detect_tokenizer_dict_param, + create_pipeline, ) @@ -38,6 +39,35 @@ def test_unknown_task_not_in_map(self) -> None: assert "image-classification" not in _HF_PIPELINE_TASK_MAP +class TestMgpstrPipeline: + def test_three_heads_are_decoded_to_generated_text(self) -> None: + import torch + + processor = MagicMock() + processor.image_processor.size = {"height": 32, "width": 128} + processor.return_value = {"pixel_values": torch.zeros((1, 3, 32, 128))} + processor.batch_decode.return_value = { + "generated_text": ["hello"], + "scores": [torch.tensor(0.9)], + } + model = MagicMock() + model.config.model_type = "mgp-str" + model.io_config = {"input_shapes": [[1, 3, 32, 128]]} + model.return_value.logits = ( + torch.zeros((1, 27, 38)), + torch.zeros((1, 27, 50257)), + torch.zeros((1, 27, 30522)), + ) + + with patch("transformers.AutoProcessor.from_pretrained", return_value=processor): + pipe = create_pipeline("image-to-text", model, "local/mgp-str") + result = pipe("image") + + processor.batch_decode.assert_called_once_with(model.return_value.logits) + assert result[0]["generated_text"] == "hello" + assert abs(result[0]["score"] - 0.9) < 1e-6 + + # --------------------------------------------------------------------------- # _detect_tokenizer_dict_param # --------------------------------------------------------------------------- diff --git a/tests/unit/loader/test_resolve_task.py b/tests/unit/loader/test_resolve_task.py index fffd20fe1..6432e2ad0 100644 --- a/tests/unit/loader/test_resolve_task.py +++ b/tests/unit/loader/test_resolve_task.py @@ -118,6 +118,20 @@ def test_unimportable_architecture_falls_back_to_hf_task_default(): assert r.source == TaskSource.HF_TASK_DEFAULT +def test_mgp_str_stale_architecture_uses_image_to_text_pipeline_tag(monkeypatch): + """Canonical family registration admits pipeline metadata after stale arch failure.""" + cfg = _cfg("mgp-str", ["MGPSTRModel"]) + cfg._name_or_path = "local-mgp-str-checkpoint" + monkeypatch.setattr( + "winml.modelkit.utils.hub_utils.get_pipeline_tag", + lambda _model_id: "image-to-text", + ) + r = resolve_task(cfg) + assert r.task == "image-to-text" + assert r.model_class.__name__ == "MgpstrForSceneTextRecognition" + assert r.source == TaskSource.PIPELINE_TAG + + # --- ported from the deleted test_detect_task_from_config.py ----------------- # These pin the architecture-class -> TasksManager task inference that the # (now-removed) ``_detect_task_from_config`` covered, asserted against the diff --git a/tests/unit/models/auto/test_auto_model.py b/tests/unit/models/auto/test_auto_model.py index ecc7d1fb2..3ffaf0c62 100644 --- a/tests/unit/models/auto/test_auto_model.py +++ b/tests/unit/models/auto/test_auto_model.py @@ -170,6 +170,53 @@ def test_get_winml_class_unsupported_task_returns_generic(self): model_class = get_winml_class("unknown", "unsupported-task-type") assert model_class == WinMLModelForGenericTask + def test_mgp_str_image_to_text_uses_specialized_wrapper(self): + from winml.modelkit.models import get_winml_class + from winml.modelkit.models.winml.image_to_text import ( + WinMLModelForMgpstrSceneTextRecognition, + ) + + assert ( + get_winml_class("mgp-str", "image-to-text") + is WinMLModelForMgpstrSceneTextRecognition + ) + + def test_mgp_str_wrapper_preserves_three_head_order(self): + from unittest.mock import MagicMock + + import torch + + from winml.modelkit.models.winml.image_to_text import ( + WinMLModelForMgpstrSceneTextRecognition, + ) + + model = object.__new__(WinMLModelForMgpstrSceneTextRecognition) + model._format_inputs = MagicMock(side_effect=lambda **kwargs: kwargs) + expected = { + "char_logits": torch.tensor([1.0]), + "bpe_logits": torch.tensor([2.0]), + "wp_logits": torch.tensor([3.0]), + } + model._run_inference = MagicMock(return_value=expected) + + result = model.forward(pixel_values=torch.zeros((1, 3, 32, 128))) + + assert result.logits == ( + expected["char_logits"], + expected["bpe_logits"], + expected["wp_logits"], + ) + + def test_mgp_str_wrapper_requires_pixel_values(self): + from winml.modelkit.models.winml.image_to_text import ( + WinMLModelForMgpstrSceneTextRecognition, + ) + + model = object.__new__(WinMLModelForMgpstrSceneTextRecognition) + + with pytest.raises(ValueError, match="requires 'pixel_values'"): + model.forward() + @pytest.mark.parametrize( "task,model_type,expected_class_name", [