diff --git a/docs/reference/index.md b/docs/reference/index.md index 086f12cf0..9927477a2 100644 --- a/docs/reference/index.md +++ b/docs/reference/index.md @@ -67,6 +67,7 @@ stages based on the target device and precision. | `dtype` | `str \| null` | Data type (e.g., `float32`, `int64`). | | `shape` | `list[int \| str] \| null` | Tensor shape (e.g., `[1, 3, 224, 224]`). String entries declare symbolic dynamic axes and use size `1` for dummy inputs. | | `value_range` | `[float, float] \| null` | Min/max for dummy tensor generation. | +| `dummy_value_runs` | `list[[int, int \| float]] \| null` | Run-length encoded deterministic dummy values. Generated configs use this only when it compactly preserves semantic input values; runs must fill the concrete tensor shape exactly. | --- diff --git a/examples/recipes/README.md b/examples/recipes/README.md index 1077e4f74..1791148b1 100644 --- a/examples/recipes/README.md +++ b/examples/recipes/README.md @@ -14,7 +14,7 @@ Each *(model, task)* includes: ## Models -Total: **75** (model, task) tuples that pass fp16 eval on all 10 (EP, device) buckets. +Total: **76** (model, task) tuples that pass fp16 eval on all 10 (EP, device) buckets. | Model | Task | |---|---| @@ -64,6 +64,7 @@ Total: **75** (model, task) tuples that pass fp16 eval on all 10 (EP, device) bu | joeddav/xlm-roberta-large-xnli | zero-shot-classification | | laion/CLIP-ViT-B-32-laion2B-s34B-b79K | feature-extraction | | mattmdjaga/segformer_b2_clothes | image-segmentation | +| microsoft/Florence-2-base | image-to-text | | microsoft/rad-dino | image-feature-extraction | | microsoft/resnet-18 | image-classification | | microsoft/resnet-50 | image-classification | diff --git a/examples/recipes/microsoft_Florence-2-base/image-to-text_fp16_config_decoder.json b/examples/recipes/microsoft_Florence-2-base/image-to-text_fp16_config_decoder.json new file mode 100644 index 000000000..897a99d64 --- /dev/null +++ b/examples/recipes/microsoft_Florence-2-base/image-to-text_fp16_config_decoder.json @@ -0,0 +1,289 @@ +{ + "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": "decoder_input_ids", + "dtype": "int32", + "shape": [ + 1, + 1 + ], + "value_range": [ + 0, + 51289 + ] + }, + { + "name": "encoder_hidden_states", + "dtype": "float32", + "shape": [ + 1, + 16, + 768 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "decoder_attention_mask", + "dtype": "int64", + "shape": [ + 1, + 1024 + ], + "dummy_value_runs": [ + [ + 1024, + 1 + ] + ] + }, + { + "name": "cache_position", + "dtype": "int64", + "shape": [ + 1 + ] + }, + { + "name": "past_0_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_0_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_1_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_1_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_2_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_2_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_3_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_3_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_4_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_4_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_5_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_5_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "logits" + }, + { + "name": "present_0_key" + }, + { + "name": "present_0_value" + }, + { + "name": "present_1_key" + }, + { + "name": "present_1_value" + }, + { + "name": "present_2_key" + }, + { + "name": "present_2_value" + }, + { + "name": "present_3_key" + }, + { + "name": "present_3_value" + }, + { + "name": "present_4_key" + }, + { + "name": "present_4_value" + }, + { + "name": "present_5_key" + }, + { + "name": "present_5_value" + } + ], + "dynamic_axes": { + "encoder_hidden_states": { + "0": "batch_size", + "1": "source_sequence_length" + } + } + }, + "optim": { + "gelu_fusion": true, + "layer_norm_fusion": true, + "matmul_add_fusion": true + }, + "quant": null, + "compile": null, + "loader": { + "task": "text2text-generation", + "model_class": "Florence2DecoderWrapper", + "model_type": "florence2", + "trust_remote_code": true + } +} \ No newline at end of file diff --git a/examples/recipes/microsoft_Florence-2-base/image-to-text_fp16_config_encoder.json b/examples/recipes/microsoft_Florence-2-base/image-to-text_fp16_config_encoder.json new file mode 100644 index 000000000..fddd038dd --- /dev/null +++ b/examples/recipes/microsoft_Florence-2-base/image-to-text_fp16_config_encoder.json @@ -0,0 +1,107 @@ +{ + "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": "int64", + "shape": [ + 1, + 8 + ], + "dummy_value_runs": [ + [ + 8, + 0 + ] + ] + }, + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 768, + 768 + ], + "dummy_value_runs": [ + [ + 1769472, + 0.0 + ] + ] + }, + { + "name": "attention_mask", + "dtype": "int64", + "shape": [ + 1, + 8 + ], + "dummy_value_runs": [ + [ + 8, + 1 + ] + ] + } + ], + "output_tensors": [ + { + "name": "last_hidden_state" + } + ], + "dynamic_axes": { + "input_ids": { + "0": "batch_size", + "1": "sequence_length" + }, + "pixel_values": { + "0": "batch_size" + }, + "attention_mask": { + "0": "batch_size", + "1": "sequence_length" + }, + "last_hidden_state": { + "0": "batch_size", + "1": "sequence_length" + } + } + }, + "optim": { + "gelu_fusion": true, + "layer_norm_fusion": true, + "matmul_add_fusion": true + }, + "quant": null, + "compile": null, + "loader": { + "task": "image-feature-extraction", + "model_class": "Florence2EncoderWrapper", + "model_type": "florence2", + "trust_remote_code": true + }, + "eval": { + "task": "image-to-text", + "prompt": "", + "dataset": { + "path": "lmms-lab/flickr30k", + "split": "test", + "samples": 500, + "columns_mapping": { + "input_column": "image", + "label_column": "caption" + } + } + } +} \ No newline at end of file diff --git a/scripts/e2e_eval/run_eval.py b/scripts/e2e_eval/run_eval.py index 8d4a86010..0500143f5 100644 --- a/scripts/e2e_eval/run_eval.py +++ b/scripts/e2e_eval/run_eval.py @@ -780,6 +780,8 @@ def _run_recipe_build( "-m", entry.hf_id, "--use-cache", + "--precision", + variant.precision, ] if no_compile: build_args += ["--no-compile"] diff --git a/scripts/e2e_eval/testsets/models_all.json b/scripts/e2e_eval/testsets/models_all.json index d63425bd7..43064d15c 100644 --- a/scripts/e2e_eval/testsets/models_all.json +++ b/scripts/e2e_eval/testsets/models_all.json @@ -5225,5 +5225,19 @@ "last_update_time": "2024-08-16T10:14:52+00:00", "optimum_supported": true, "order": 5 + }, + { + "hf_id": "microsoft/Florence-2-base", + "task": "image-to-text", + "model_type": "florence2", + "group": "microsoft", + "priority": "P1", + "downloads": 2651904, + "last_update_time": "2025-08-04T17:36:57+00:00", + "optimum_supported": false, + "order": 1, + "tags": [ + "acc" + ] } ] diff --git a/scripts/e2e_eval/testsets/models_with_acc.json b/scripts/e2e_eval/testsets/models_with_acc.json index e21423161..550d02033 100644 --- a/scripts/e2e_eval/testsets/models_with_acc.json +++ b/scripts/e2e_eval/testsets/models_with_acc.json @@ -2064,5 +2064,24 @@ "mask_column": "mask" } } + }, + { + "hf_id": "microsoft/Florence-2-base", + "task": "image-to-text", + "model_type": "florence2", + "group": "microsoft", + "priority": "P1", + "precision": "fp16", + "dataset_config": { + "path": "lmms-lab/flickr30k", + "split": "test", + "samples": 500, + "metric": "cider", + "winml_metric_key": "cider", + "columns_mapping": { + "input_column": "image", + "label_column": "caption" + } + } } ] diff --git a/src/winml/modelkit/analyze/core/runtime_checker_query.py b/src/winml/modelkit/analyze/core/runtime_checker_query.py index 18382f57d..a53618b16 100644 --- a/src/winml/modelkit/analyze/core/runtime_checker_query.py +++ b/src/winml/modelkit/analyze/core/runtime_checker_query.py @@ -25,6 +25,7 @@ ONNXDomain, SupportedONNXType, infer_onnx_shapes, + infer_symbolic_shapes, remove_optional_from_type_annotation, ) from ...onnx.external_data import try_load_external_initializer_array @@ -1018,9 +1019,7 @@ def __init__( # Then try to enhance with symbolic shape inference # if available which supports Microsoft domain try: - from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference - - symbolic_inferred = SymbolicShapeInference.infer_shapes(inferred_model) + symbolic_inferred = infer_symbolic_shapes(inferred_model) if symbolic_inferred is not None: inferred_model = symbolic_inferred except Exception as e: diff --git a/src/winml/modelkit/commands/inspect.py b/src/winml/modelkit/commands/inspect.py index 579e6c92d..0a01d0460 100644 --- a/src/winml/modelkit/commands/inspect.py +++ b/src/winml/modelkit/commands/inspect.py @@ -146,6 +146,12 @@ def _list_tasks_for_model(model_type: str) -> list[str]: default=None, help="Override model class (e.g., BertForMaskedLM) — can be used without --model", ) +@click.option( + "--trust-remote-code", + is_flag=True, + default=False, + help="Trust remote/custom HuggingFace code when loading model configuration", +) @cli_utils.verbosity_options() @cli_utils.no_color_option() @click.pass_context @@ -160,6 +166,7 @@ def inspect( list_tasks: bool, model_type: str | None, model_class: str | None, + trust_remote_code: bool, ) -> None: r"""Inspect input model's WinML CLI configuration. @@ -196,7 +203,10 @@ def inspect( from transformers import AutoConfig try: - hf_config = AutoConfig.from_pretrained(model, trust_remote_code=False) + hf_config = AutoConfig.from_pretrained( + model, + trust_remote_code=trust_remote_code, + ) except Exception as e: raise click.ClickException( f"Could not resolve model type for '{model}': {e}" @@ -299,6 +309,7 @@ def inspect( model_type_override=model_type, model_class_override=model_class, include_hierarchy=hierarchy, + trust_remote_code=trust_remote_code, ) else: with _stderr_console.status( @@ -311,6 +322,7 @@ def inspect( model_type_override=model_type, model_class_override=model_class, include_hierarchy=hierarchy, + trust_remote_code=trust_remote_code, ) if output_format == "json": @@ -338,6 +350,7 @@ def _inspect_model_v2( model_type_override: str | None = None, model_class_override: str | None = None, include_hierarchy: bool = False, + trust_remote_code: bool = False, ) -> InspectResult: """Inspect v2 core — calls shared loader/export modules directly. @@ -347,6 +360,7 @@ def _inspect_model_v2( model_type_override: Model type override (e.g., "bert") model_class_override: Model class override (e.g., "BertForMaskedLM") include_hierarchy: Whether to extract module hierarchy + trust_remote_code: Whether to trust remote/custom HuggingFace code Returns: InspectResult dataclass @@ -368,6 +382,7 @@ def _inspect_model_v2( build_tensor_infos_from_io_specs, compile_support_status, resolve_cache, + resolve_composite_exporter, resolve_composite_info, resolve_io_config, resolve_processor, @@ -389,7 +404,10 @@ def _inspect_model_v2( parent_hf_config = None if model_id and not model_type_override: try: - parent_hf_config = AutoConfig.from_pretrained(model_id, trust_remote_code=False) + parent_hf_config = AutoConfig.from_pretrained( + model_id, + trust_remote_code=trust_remote_code, + ) except Exception: pass # resolve_loader_config will handle the error properly @@ -405,6 +423,7 @@ def _inspect_model_v2( model_type=model_type_override, model_class=model_class_override, hf_config=parent_hf_config, + trust_remote_code=trust_remote_code, ) except RepositoryNotFoundError as e: # Direct HF Hub 404 — keep full message (includes private-repo hint). @@ -531,6 +550,14 @@ def _inspect_model_v2( output_tensors=output_tensors, opset_version=opset_version, ) + composite_exporter = resolve_composite_exporter( + model_type, + task, + hf_config=hf_config, + model_id=model_id, + ) + if composite_exporter is not None: + exporter_info = composite_exporter # ========================================================================= # STEP 6: WinML class (inspect-only lookup) @@ -545,7 +572,7 @@ def _inspect_model_v2( try: from ..inspect.hierarchy import extract_hierarchy - hierarchy_info = extract_hierarchy(model_id) + hierarchy_info = extract_hierarchy(model_id, trust_remote_code=trust_remote_code) except Exception as e: logger.debug("Hierarchy extraction failed for %s: %s", model_id, e) diff --git a/src/winml/modelkit/eval/base_evaluator.py b/src/winml/modelkit/eval/base_evaluator.py index 2e689b13f..67466da8e 100644 --- a/src/winml/modelkit/eval/base_evaluator.py +++ b/src/winml/modelkit/eval/base_evaluator.py @@ -131,27 +131,10 @@ def prepare_data(self) -> Dataset: def prepare_pipeline(self) -> Pipeline: """Create HF pipeline for inference. Subclasses override to configure.""" - from transformers import pipeline + from ..inference.pipeline import create_pipeline assert self.config.task is not None, "config.task is required to build pipeline" - pipeline_task = _PIPELINE_TASK_MAP.get(self.config.task, self.config.task) - # transformers.pipeline has 60+ Literal overloads — runtime task strings - # can't be statically matched. The string-task fallback handles unknown tasks. - return cast( - "Pipeline", - pipeline( # type: ignore[call-overload, misc] # 60+ Literal overloads + union model arg - pipeline_task, - model=self.model, - framework="pt", - tokenizer=self.config.model_id, - feature_extractor=self.config.model_id, - image_processor=self.config.model_id, - processor=self.config.model_id, - # "device" is for HF pipeline pytorch tensors, not ORT EP. - # WinMLSession handles device delegation for ORT. - device="cpu", - ), - ) + return cast("Pipeline", create_pipeline(self.config.task, self.model, self.config.model_id)) def _fixed_seq_length(self) -> int | None: """Return the model's fixed sequence length, or ``None`` if dynamic. diff --git a/src/winml/modelkit/eval/config.py b/src/winml/modelkit/eval/config.py index 923c0474d..8054be52e 100644 --- a/src/winml/modelkit/eval/config.py +++ b/src/winml/modelkit/eval/config.py @@ -89,6 +89,7 @@ class WinMLEvaluationConfig: composite models (e.g. ``{"image-encoder": "...", "text-encoder": "..."}``). None = build from model_id. task: HF pipeline task. Auto-detected from model_id if omitted. + prompt: Optional task prompt passed to prompt-aware pipelines. device: Target device for inference. ep: Explicit execution provider (e.g., "qnn", "dml"). Overrides device-to-provider mapping when provided. @@ -112,6 +113,7 @@ class WinMLEvaluationConfig: model_id: str | None = None model_path: str | dict[str, str] | None = None task: str | None = None + prompt: str | None = None device: str = "auto" precision: str = "auto" ep: EPNameOrAlias | None = None @@ -136,6 +138,8 @@ def to_dict(self) -> dict: result["model_path"] = self.model_path if self.task is not None: result["task"] = self.task + if self.prompt is not None: + result["prompt"] = self.prompt result["device"] = self.device if self.precision != "auto": result["precision"] = self.precision @@ -182,6 +186,7 @@ def from_dict(cls, data: dict) -> WinMLEvaluationConfig: model_id=data.get("model_id"), model_path=data.get("model_path"), task=data.get("task"), + prompt=data.get("prompt"), device=data.get("device", "auto"), precision=data.get("precision", "auto"), ep=data.get("ep"), diff --git a/src/winml/modelkit/eval/image_to_text_evaluator.py b/src/winml/modelkit/eval/image_to_text_evaluator.py index 9f541f173..027f50bf3 100644 --- a/src/winml/modelkit/eval/image_to_text_evaluator.py +++ b/src/winml/modelkit/eval/image_to_text_evaluator.py @@ -73,7 +73,8 @@ def compute(self) -> dict[str, Any]: continue try: - out = self.pipe(image) + kwargs = {"prompt": self.config.prompt} if self.config.prompt is not None else {} + out = self.pipe(image, **kwargs) except Exception as e: logger.warning("Pipeline call failed (skipping): %s", e) skipped += 1 diff --git a/src/winml/modelkit/export/config.py b/src/winml/modelkit/export/config.py index 34ac3ba8c..a16d856aa 100644 --- a/src/winml/modelkit/export/config.py +++ b/src/winml/modelkit/export/config.py @@ -477,6 +477,7 @@ def _resolve_export_config_from_specs( ) value_ranges = io_specs.get("value_ranges", {}) + dummy_value_runs = io_specs.get("dummy_value_runs", {}) input_tensors = [ InputTensorSpec( @@ -484,6 +485,7 @@ def _resolve_export_config_from_specs( shape=shape, dtype=dtype, value_range=value_ranges.get(name), + dummy_value_runs=dummy_value_runs.get(name), ) for name, shape, dtype in zip(input_names, input_shapes, input_dtypes, strict=False) ] diff --git a/src/winml/modelkit/export/io.py b/src/winml/modelkit/export/io.py index 3bedca226..95976dbc8 100644 --- a/src/winml/modelkit/export/io.py +++ b/src/winml/modelkit/export/io.py @@ -41,6 +41,7 @@ ) from ..loader import to_optimum_task +from ..onnx import InputTensorSpec from .value_range import intercept_value_ranges @@ -486,13 +487,20 @@ def resolve_io_specs( input_shapes = [tuple(t.shape) for t in dummy_inputs.values()] input_dtypes = [str(t.dtype).replace("torch.", "") for t in dummy_inputs.values()] + dummy_value_runs = {} + if getattr(onnx_config, "PRESERVE_DUMMY_VALUE_RUNS", False): + dummy_value_runs = { + name: value_runs + for name, tensor in dummy_inputs.items() + if (value_runs := InputTensorSpec.compact_dummy_value_runs(tensor)) is not None + } # Build value_range dict: {name: (min, max)} from intercepted data value_range_tuples = { name: (info["min"], info["max"]) for name, info in value_ranges.items() } - return { + specs = { "inputs": onnx_config.inputs, "outputs": onnx_config.outputs, "input_names": list(onnx_config.inputs.keys()), @@ -502,3 +510,6 @@ def resolve_io_specs( "input_dtypes": input_dtypes, "value_ranges": value_range_tuples, } + if dummy_value_runs: + specs["dummy_value_runs"] = dummy_value_runs + return specs diff --git a/src/winml/modelkit/inference/pipeline.py b/src/winml/modelkit/inference/pipeline.py index fc8bdb66c..89c5e9b08 100644 --- a/src/winml/modelkit/inference/pipeline.py +++ b/src/winml/modelkit/inference/pipeline.py @@ -20,12 +20,18 @@ import inspect import logging -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Protocol, cast + +from transformers.pipelines.image_to_text import ImageToTextPipeline + +from ..models.winml.composite_model import PipelineCapability if TYPE_CHECKING: from collections.abc import Mapping + from transformers.pipelines.base import GenericTensor + from ..models.winml.base import WinMLPreTrainedModel from ..models.winml.composite_model import WinMLCompositeModel @@ -38,6 +44,122 @@ } +class SupportsPipelineCapabilities(Protocol): + """Model protocol for selecting non-default preprocessing pipelines.""" + + pipeline_capabilities: frozenset[PipelineCapability] + + +class SupportsCombinedProcessor(SupportsPipelineCapabilities, Protocol): + """Model protocol for combined image/text processor construction.""" + + def create_combined_processor(self, model_id: str) -> Any: + """Load the processor that satisfies the model's declared contract.""" + + +class SupportsTokenDecoding(Protocol): + """Tokenizer capability required by image-to-text postprocessing.""" + + def decode(self, token_ids: Any, *, skip_special_tokens: bool) -> str: + """Decode generated token IDs.""" + + +class SupportsTokenizer(Protocol): + """Processor capability for supplying the postprocessing tokenizer.""" + + tokenizer: SupportsTokenDecoding + + +class SupportsCombinedProcessorInputs(Protocol): + """Processor output that can be transferred to the pipeline tensor dtype.""" + + def to(self, device: object) -> SupportsCombinedProcessorInputs: + """Move the processor output to a tensor device or dtype.""" + + +class SupportsCombinedImageTextProcessor(SupportsTokenizer, Protocol): + """Combined image/text processor surface required by the custom pipeline.""" + + def __call__( + self, *, images: object, text: str, return_tensors: str + ) -> SupportsCombinedProcessorInputs: + """Process an image and its text prompt together.""" + + +class CombinedProcessorImageToTextPipeline(ImageToTextPipeline): + """Image-to-text pipeline that preserves a processor's joint image/text contract.""" + + _load_processor = True + _load_image_processor = False + _load_feature_extractor = False + _load_tokenizer = False + + # Transformers' Pipeline stub uses ``input_`` plus ``**dict`` while its + # ImageToTextPipeline override uses image/prompt/timeout. Preserve the + # latter's public API and narrow only the incompatible base-stub override. + def preprocess( # type: ignore[override] + self, image: Any, prompt: Any = None, timeout: Any = None + ) -> dict[str, GenericTensor]: + """Create model inputs with one combined processor invocation.""" + from transformers.image_utils import load_image + + if prompt is None: + raise ValueError("A prompt is required by the combined image/text processor.") + processor = self.processor + if processor is None or not callable(processor): + raise TypeError("A combined image/text processor is required.") + image = load_image(image, timeout=timeout) + model_inputs = cast("SupportsCombinedImageTextProcessor", processor)( + images=image, + text=prompt, + return_tensors=self.framework, + ) + if self.framework == "pt": + model_inputs = model_inputs.to(self.dtype) + return cast("dict[str, GenericTensor]", model_inputs) + + +def _pipeline_class_for(model: Any) -> type | None: + """Resolve an HF pipeline implementation from declared model capabilities.""" + capabilities = inspect.getattr_static(model, "pipeline_capabilities", frozenset()) + if not isinstance(capabilities, frozenset): + raise TypeError("pipeline_capabilities must be a frozenset of PipelineCapability values") + if not all(isinstance(capability, PipelineCapability) for capability in capabilities): + raise TypeError("pipeline_capabilities must contain PipelineCapability values") + if PipelineCapability.COMBINED_IMAGE_TEXT_PROCESSOR in capabilities: + return CombinedProcessorImageToTextPipeline + return None + + +def _combined_processor_for( + model: Any, model_id: str | None +) -> SupportsCombinedImageTextProcessor: + """Load the declared combined processor with explicit capability errors.""" + loader = getattr(model, "create_combined_processor", None) + if not callable(loader): + raise TypeError( + "Models declaring combined-image-text-processor must implement " + "create_combined_processor(model_id)." + ) + if model_id is None: + raise ValueError( + "A model ID is required to load a combined image/text processor." + ) + processor = loader(model_id) + tokenizer = getattr(processor, "tokenizer", None) + if not callable(processor) or not callable(getattr(tokenizer, "decode", None)): + raise TypeError( + "Combined image/text processors must be callable and expose a tokenizer " + "with a decode method." + ) + return cast("SupportsCombinedImageTextProcessor", processor) + + +def _tokenizer_for(processor: SupportsTokenizer) -> SupportsTokenDecoding: + """Return the processor-owned tokenizer required by image-to-text decoding.""" + return processor.tokenizer + + def create_pipeline( task: str, model: WinMLPreTrainedModel | WinMLCompositeModel, @@ -65,7 +187,13 @@ def create_pipeline( # WinMLSession handles device delegation internally. "device": "cpu", } - if model_id: + pipeline_class = _pipeline_class_for(model) + if pipeline_class is not None: + processor = _combined_processor_for(model, model_id) + kwargs["pipeline_class"] = pipeline_class + kwargs["processor"] = processor + kwargs["tokenizer"] = _tokenizer_for(processor) + elif model_id: kwargs["tokenizer"] = model_id kwargs["feature_extractor"] = model_id kwargs["image_processor"] = model_id @@ -129,7 +257,10 @@ def _adapt_tokenizer_padding(pipe: Any, task: str, model: Any) -> None: # No **kwargs — only accepts specific named params # → set only params that appear in the signature - preprocess_sig = inspect.signature(type(pipe).preprocess) + preprocess = getattr(type(pipe), "preprocess", None) + if not callable(preprocess): + return + preprocess_sig = inspect.signature(preprocess) sig_params = preprocess_sig.parameters tok_dict_key = _detect_tokenizer_dict_param(pipe, sig_params) diff --git a/src/winml/modelkit/inspect/__init__.py b/src/winml/modelkit/inspect/__init__.py index 318c2aa7b..dd06492c8 100644 --- a/src/winml/modelkit/inspect/__init__.py +++ b/src/winml/modelkit/inspect/__init__.py @@ -28,6 +28,7 @@ get_build_config, get_known_tasks, resolve_cache, + resolve_composite_exporter, resolve_composite_info, resolve_exporter, resolve_io_config, @@ -72,6 +73,7 @@ def inspect_model( model_id: str, include_hierarchy: bool = False, task_override: str | None = None, + trust_remote_code: bool = False, ) -> InspectResult: """Inspect a HuggingFace model and return configuration details. @@ -79,6 +81,7 @@ def inspect_model( model_id: HuggingFace model identifier (e.g., "openai/clip-vit-base-patch32") include_hierarchy: If True, load model and extract HF module hierarchy task_override: If provided, use this task instead of auto-detection + trust_remote_code: Whether to trust remote/custom code when loading config Returns: InspectResult with all configuration details @@ -98,7 +101,7 @@ def inspect_model( # Step 1: Fetch HF config (no model download) try: - hf_config = AutoConfig.from_pretrained(model_id, trust_remote_code=False) + hf_config = AutoConfig.from_pretrained(model_id, trust_remote_code=trust_remote_code) except OSError as e: if "404" in str(e) or "not found" in str(e).lower(): raise ModelNotFoundError(f"Model '{model_id}' not found on HuggingFace Hub") from e @@ -136,7 +139,12 @@ def inspect_model( ) # Step 4: Resolve exporter configuration (pass model_id for correct image sizes) - exporter_info = resolve_exporter(model_type, task, hf_config=hf_config, model_id=model_id) + exporter_info = resolve_composite_exporter( + model_type, + task, + hf_config=hf_config, + model_id=model_id, + ) or resolve_exporter(model_type, task, hf_config=hf_config, model_id=model_id) logger.debug( "Exporter: %s (source: %s)", exporter_info.onnx_config_class, @@ -152,7 +160,7 @@ def inspect_model( if include_hierarchy: from .hierarchy import extract_hierarchy - hierarchy_info = extract_hierarchy(model_id) + hierarchy_info = extract_hierarchy(model_id, trust_remote_code=trust_remote_code) logger.debug("Hierarchy: %d HF modules", hierarchy_info.hf_module_count) # Step 6: Compile overall support status @@ -233,6 +241,7 @@ def inspect_model( "get_known_tasks", "inspect_model", "resolve_cache", + "resolve_composite_exporter", "resolve_composite_info", "resolve_io_config", "resolve_processor", diff --git a/src/winml/modelkit/inspect/hierarchy.py b/src/winml/modelkit/inspect/hierarchy.py index e375e9073..b91eb66e6 100644 --- a/src/winml/modelkit/inspect/hierarchy.py +++ b/src/winml/modelkit/inspect/hierarchy.py @@ -129,7 +129,7 @@ def _is_hf_module(module: nn.Module) -> bool: return True -def extract_hierarchy(model_id: str) -> HierarchyInfo: +def extract_hierarchy(model_id: str, trust_remote_code: bool = False) -> HierarchyInfo: """Extract the HF module hierarchy from a model. If the model is already cached locally, loads pretrained weights. @@ -138,6 +138,7 @@ def extract_hierarchy(model_id: str) -> HierarchyInfo: Args: model_id: HuggingFace model identifier + trust_remote_code: Whether to trust remote/custom HuggingFace code. Returns: HierarchyInfo with the module hierarchy @@ -149,15 +150,15 @@ def extract_hierarchy(model_id: str) -> HierarchyInfo: logger.debug("Checking if model is cached locally...") model = AutoModel.from_pretrained( model_id, - trust_remote_code=False, + trust_remote_code=trust_remote_code, local_files_only=True, ) logger.info("Using cached pretrained model") except Exception as e: # Model not cached locally or other loading issue - use random weights logger.debug("Model not cached or load failed (%s), using random weights", type(e).__name__) - config = AutoConfig.from_pretrained(model_id, trust_remote_code=False) - model = AutoModel.from_config(config) + config = AutoConfig.from_pretrained(model_id, trust_remote_code=trust_remote_code) + model = AutoModel.from_config(config, trust_remote_code=trust_remote_code) logger.info("Using random weights (model not downloaded)") model.eval() diff --git a/src/winml/modelkit/inspect/resolver.py b/src/winml/modelkit/inspect/resolver.py index c5a6e2c01..c692a82c1 100644 --- a/src/winml/modelkit/inspect/resolver.py +++ b/src/winml/modelkit/inspect/resolver.py @@ -270,17 +270,9 @@ def resolve_exporter( model_type_normalized = model_type.lower().replace("_", "-") # Check MODEL_BUILD_CONFIGS for predefined config - if model_type_normalized in MODEL_BUILD_CONFIGS: - config: WinMLBuildConfig = MODEL_BUILD_CONFIGS[model_type_normalized] - # MODEL_BUILD_CONFIGS entries are HF export configs; export is None only on - # the direct-ONNX build path, which never reaches this registry lookup. - export_config = config.export - if export_config is None: - raise ValueError( - f"MODEL_BUILD_CONFIGS entry for {model_type_normalized!r} is missing an " - "export config (export is None only on the direct-ONNX build path)." - ) - + config: WinMLBuildConfig | None = MODEL_BUILD_CONFIGS.get(model_type_normalized) + export_config = config.export if config is not None else None + if export_config is not None and export_config.input_tensors is not None: # Extract input tensors input_tensors: list[TensorInfo] = [] if export_config.input_tensors: @@ -375,6 +367,47 @@ def resolve_exporter( ) +def resolve_composite_exporter( + model_type: str, + task: str, + hf_config: PretrainedConfig | None = None, + *, + model_id: str | None = None, +) -> ExporterInfo | None: + """Resolve a composite exporter by validating every registered component.""" + from ..models.winml.composite_model import COMPOSITE_MODEL_REGISTRY + + model_type_normalized = model_type.lower().replace("_", "-") + composite_cls = COMPOSITE_MODEL_REGISTRY.get((model_type_normalized, task)) + if composite_cls is None: + return None + + component_names = list(composite_cls._SUB_MODEL_CONFIG) + for component_task in composite_cls._SUB_MODEL_CONFIG.values(): + loader = resolve_loader(model_type_normalized, component_task) + exporter = resolve_exporter( + model_type_normalized, + component_task, + hf_config=hf_config, + model_id=model_id, + ) + if ( + loader.support_level == SupportLevel.UNSUPPORTED + or exporter.support_level == SupportLevel.UNSUPPORTED + ): + return ExporterInfo( + onnx_config_class=None, + onnx_config_source="COMPOSITE_MODEL_REGISTRY", + support_level=SupportLevel.UNSUPPORTED, + ) + + return ExporterInfo( + onnx_config_class=f"Composite ({', '.join(component_names)})", + onnx_config_source="COMPOSITE_MODEL_REGISTRY", + support_level=SupportLevel.SUPPORTED, + ) + + def resolve_winml(model_type: str, task: str) -> WinMLInfo: """Resolve WinML inference class for a model. @@ -392,6 +425,16 @@ def resolve_winml(model_type: str, task: str) -> WinMLInfo: """ model_type_normalized = model_type.lower().replace("_", "-") + from ..models.winml.composite_model import COMPOSITE_MODEL_REGISTRY + + composite_cls = COMPOSITE_MODEL_REGISTRY.get((model_type_normalized, task)) + if composite_cls is not None: + return WinMLInfo( + winml_class=composite_cls.__name__, + winml_class_source="COMPOSITE_MODEL_REGISTRY", + support_level=SupportLevel.SUPPORTED, + ) + # Level 1: Check WINML_MODEL_CLASS_MAPPING (specialized) key = (model_type_normalized, task) if key in WINML_MODEL_CLASS_MAPPING: diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index 094714494..a03256f9c 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/src/winml/modelkit/models/hf/__init__.py @@ -44,6 +44,14 @@ from .depth_anything import DepthAnythingIOConfig as _DepthAnythingIOConfig # triggers registration from .depth_pro import DepthProIOConfig as _DepthProIOConfig # triggers registration from .detr import DETR_CONFIG +from .florence2 import FLORENCE2_CONFIG +from .florence2 import MODEL_CLASS_MAPPING as _FLORENCE2_CLASS_MAPPING +from .florence2 import ( + Florence2DecoderIOConfig as _Florence2DecoderIOConfig, # triggers registration +) +from .florence2 import ( + Florence2EncoderIOConfig as _Florence2EncoderIOConfig, # 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 @@ -116,6 +124,7 @@ _BART_CLASS_MAPPING, _BLIP_CLASS_MAPPING, _CLIP_CLASS_MAPPING, + _FLORENCE2_CLASS_MAPPING, _MARIAN_CLASS_MAPPING, _MU2_CLASS_MAPPING, _QWEN_CLASS_MAPPING, @@ -144,6 +153,7 @@ "clip-text-model": CLIP_CONFIG, "clip-vision-model": CLIP_CONFIG, "detr": DETR_CONFIG, + "florence2": FLORENCE2_CONFIG, "marian": MARIAN_CONFIG, "roberta": ROBERTA_FAMILY_CONFIG, "mu2": MU2_CONFIG, diff --git a/src/winml/modelkit/models/hf/decoder_wrapper.py b/src/winml/modelkit/models/hf/decoder_wrapper.py index e83148268..bb8985a67 100644 --- a/src/winml/modelkit/models/hf/decoder_wrapper.py +++ b/src/winml/modelkit/models/hf/decoder_wrapper.py @@ -35,16 +35,29 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar, Protocol import torch import torch.nn as nn from optimum.exporters.onnx import OnnxConfig -from transformers import PreTrainedModel from ..winml.kv_cache import WinMLCache, WinMLStaticCache +if TYPE_CHECKING: + from transformers import PreTrainedModel + + +class _PreTrainedModelLoader(Protocol): + """Class-level loader contract shared by built-in and upstream model classes.""" + + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path: str, **kwargs: Any + ) -> PreTrainedModel: + """Load a pretrained model.""" + + class WinMLDecoderWrapper(nn.Module, ABC): """Abstract base class for static-KV-cache decoder export wrappers. @@ -60,7 +73,8 @@ class WinMLDecoderWrapper(nn.Module, ABC): num_layers — derived from ``onnx_config._normalized_config.num_layers`` """ - _HF_MODEL_CLS: ClassVar[type[PreTrainedModel]] # set per-subclass to a concrete HF model class + # Set per-subclass to a concrete HF model loader. + _HF_MODEL_CLS: ClassVar[_PreTrainedModelLoader] _IO_CONFIG_CLS: ClassVar[type] _TASK: ClassVar[str] = "text2text-generation" _CACHE_CLS: ClassVar[type[WinMLCache]] = WinMLStaticCache diff --git a/src/winml/modelkit/models/hf/florence2.py b/src/winml/modelkit/models/hf/florence2.py new file mode 100644 index 000000000..2a40e479f --- /dev/null +++ b/src/winml/modelkit/models/hf/florence2.py @@ -0,0 +1,423 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Florence-2 split image-to-text export.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, ClassVar, Protocol, TypedDict, cast + +import torch +import torch.nn as nn +from optimum.exporters.onnx import OnnxConfig +from optimum.utils import NormalizedConfig +from optimum.utils.input_generators import DummyInputGenerator + +from ...config import WinMLBuildConfig +from ...export import register_onnx_overwrite +from ...optim import WinMLOptimizationConfig +from ..winml.composite_model import PipelineCapability, register_composite_model +from ..winml.encoder_decoder import EncoderDecoderInputGenerator, WinMLEncoderDecoderModel +from ..winml.kv_cache import PastKeyValueInputGenerator, WinMLStaticCache +from .decoder_wrapper import WinMLDecoderWrapper, WinMLStaticCacheDecoderIOConfig + + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from transformers import GenerationConfig, PretrainedConfig, PreTrainedModel + + +FLORENCE2_CONFIG = WinMLBuildConfig( + optim=WinMLOptimizationConfig( + gelu_fusion=True, + layer_norm_fusion=True, + matmul_add_fusion=True, + ), +) + + +class _LoadingInfo(TypedDict): + """Subset of Transformers' checkpoint reconciliation metadata.""" + + missing_keys: Sequence[object] + unexpected_keys: Sequence[object] + mismatched_keys: Sequence[object] + + +class _Florence2EncoderOutput(Protocol): + """Encoder result needed by the ONNX export wrapper.""" + + last_hidden_state: torch.Tensor + + +class _Florence2Encoder(Protocol): + """Callable encoder surface exposed by Florence-2.""" + + def __call__( + self, + *, + attention_mask: torch.Tensor, + inputs_embeds: torch.Tensor, + return_dict: bool, + ) -> _Florence2EncoderOutput: + """Encode the supplied embeddings.""" + + +class _Florence2DecoderOutput(Protocol): + """Decoder result needed by the static-cache export wrapper.""" + + last_hidden_state: torch.Tensor + past_key_values: tuple[tuple[torch.Tensor, ...], ...] + + +class _Florence2Decoder(Protocol): + """Callable decoder surface exposed by Florence-2.""" + + def __call__( + self, + *, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: None, + past_key_values: tuple[tuple[torch.Tensor, torch.Tensor], ...], + use_cache: bool, + return_dict: bool, + ) -> _Florence2DecoderOutput: + """Decode a token using the supplied encoder states and cache.""" + + +class _Florence2LanguageModel(Protocol): + """Language-model head surface used after Florence-2 decoding.""" + + lm_head: Callable[[torch.Tensor], torch.Tensor] + + +class _NativeFlorence2ForConditionalGeneration: + """Load the checkpoint with its upstream model implementation.""" + + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path: str, **kwargs: Any + ) -> PreTrainedModel: + """Load all checkpoint tensors through the model's upstream contract.""" + from transformers import AutoModelForCausalLM + + kwargs["output_loading_info"] = True + kwargs.setdefault("attn_implementation", "eager") + model, loading_info = cast( + "tuple[PreTrainedModel, _LoadingInfo]", + AutoModelForCausalLM.from_pretrained( + pretrained_model_name_or_path, **kwargs + ), + ) + unresolved = { + name: keys + for name, keys in ( + ("missing_keys", loading_info["missing_keys"]), + ("unexpected_keys", loading_info["unexpected_keys"]), + ("mismatched_keys", loading_info["mismatched_keys"]), + ) + if keys + } + if unresolved: + counts = ", ".join(f"{name}={len(keys)}" for name, keys in unresolved.items()) + raise RuntimeError(f"Checkpoint reconciliation failed: {counts}.") + return model + + +def _load_native_combined_processor(model_id: str, *, trust_remote_code: bool = False) -> Any: + """Load the processor that accompanies the upstream model implementation.""" + from transformers import AutoProcessor + + return AutoProcessor.from_pretrained(model_id, trust_remote_code=trust_remote_code) + + +class _Florence2EncoderInputGenerator(DummyInputGenerator): # type: ignore[misc] # optimum base is untyped + """Generate image placeholders together with a caption prompt.""" + + SUPPORTED_INPUT_NAMES = ("input_ids", "pixel_values", "attention_mask") + + def __init__(self, task: str, normalized_config: Any, **kwargs: Any) -> None: + self.batch_size = kwargs.get("batch_size", 1) + self.image_size = normalized_config.image_size + self.num_channels = normalized_config.num_channels + + def generate( + self, + input_name: str, + framework: str = "pt", + int_dtype: str = "int64", + float_dtype: str = "fp32", + ) -> torch.Tensor: + del framework, int_dtype, float_dtype + sequence_length = 8 + if input_name == "input_ids": + return torch.zeros((self.batch_size, sequence_length), dtype=torch.long) + if input_name == "pixel_values": + return torch.zeros( + (self.batch_size, self.num_channels, self.image_size, self.image_size), + dtype=torch.float32, + ) + if input_name == "attention_mask": + return torch.ones((self.batch_size, sequence_length), dtype=torch.long) + raise ValueError(f"Unknown input: {input_name}") + + +class _Florence2EncoderNormalizedConfig(NormalizedConfig): # type: ignore[misc] # optimum base is untyped + """Normalize Florence-2's image and prompt dimensions.""" + + def __init__(self, config: Any, **kwargs: Any) -> None: + super().__init__(config, **kwargs) + vision_config = config.vision_config + self.num_channels = getattr( + vision_config, "num_channels", getattr(vision_config, "in_channels", 3) + ) + self.image_size = 768 + + +class Florence2EncoderWrapper(nn.Module): + """Export Florence-2's image-aware text encoder.""" + + def __init__(self, model: nn.Module, config: Any) -> None: + super().__init__() + self.model = model + self.config = config + + @classmethod + def from_pretrained(cls, model_name_or_path: str, **kwargs: Any) -> Florence2EncoderWrapper: + """Load Florence-2 and retain its image-aware encoder.""" + full = _NativeFlorence2ForConditionalGeneration.from_pretrained( + model_name_or_path, **kwargs + ) + wrapper = cls(full, full.config) + wrapper.eval() + return wrapper + + def forward( + self, + input_ids: torch.Tensor, + pixel_values: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + """Encode the caption prompt after replacing image placeholders.""" + get_input_embeddings = cast( + "Callable[[], Callable[[torch.Tensor], torch.Tensor]]", + self.model.get_input_embeddings, + ) + input_embeddings = get_input_embeddings() + inputs_embeds = input_embeddings(input_ids) + encode_image = cast("Callable[[torch.Tensor], torch.Tensor]", self.model._encode_image) + image_features = encode_image(pixel_values).to( + inputs_embeds.device, inputs_embeds.dtype + ) + merge_image_features = cast( + "Callable[[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]", + self.model._merge_input_ids_with_image_features, + ) + inputs_embeds, _ = merge_image_features(image_features, inputs_embeds) + image_attention_mask = attention_mask.new_ones( + (attention_mask.size(0), image_features.size(1)) + ) + attention_mask = torch.cat((image_attention_mask, attention_mask), dim=1) + get_encoder = cast("Callable[[], _Florence2Encoder]", self.model.get_encoder) + encoder = get_encoder() + outputs = encoder( + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + return_dict=True, + ) + return outputs.last_hidden_state + + +@register_onnx_overwrite("florence2", "feature-extraction", library_name="transformers") +class Florence2EncoderIOConfig(OnnxConfig): # type: ignore[misc] # optimum base is untyped + """ONNX config for Florence-2's prompt-aware encoder.""" + + NORMALIZED_CONFIG_CLASS = _Florence2EncoderNormalizedConfig + DUMMY_INPUT_GENERATOR_CLASSES = (_Florence2EncoderInputGenerator,) + PRESERVE_DUMMY_VALUE_RUNS = True + + @property + def inputs(self) -> dict[str, dict[int, str]]: # noqa: D102 + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: # noqa: D102 + return {"last_hidden_state": {0: "batch_size", 1: "sequence_length"}} + + +class _Florence2DecoderNormalizedConfig(NormalizedConfig): # type: ignore[misc] # optimum base is untyped + """Normalize Florence-2's BART decoder configuration.""" + + def __init__(self, config: Any, **kwargs: Any) -> None: + super().__init__(config, **kwargs) + self._text_config = config.text_config + + @property + def hidden_size(self) -> int: + return cast("int", self._text_config.d_model) + + @property + def num_layers(self) -> int: + return cast("int", self._text_config.decoder_layers) + + @property + def num_attention_heads(self) -> int: + return cast("int", self._text_config.decoder_attention_heads) + + @property + def head_dim(self) -> int: + return self.hidden_size // self.num_attention_heads + + @property + def max_cache_len(self) -> int: + return cast("int", self._text_config.max_position_embeddings) + + @property + def vocab_size(self) -> int: + return cast("int", self._text_config.vocab_size) + + +@register_onnx_overwrite("florence2", "text2text-generation", library_name="transformers") +class Florence2DecoderIOConfig(WinMLStaticCacheDecoderIOConfig): + """ONNX config for Florence-2's static-cache BART decoder.""" + + NORMALIZED_CONFIG_CLASS = _Florence2DecoderNormalizedConfig + DUMMY_INPUT_GENERATOR_CLASSES = ( + EncoderDecoderInputGenerator, + PastKeyValueInputGenerator, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: # noqa: D102 + result: dict[str, dict[int, str]] = { + "decoder_input_ids": {0: "batch_size"}, + "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, + "decoder_attention_mask": {0: "batch_size"}, + "cache_position": {}, + } + for i in range(self._normalized_config.num_layers): + result[f"past_{i}_key"] = {0: "batch_size"} + result[f"past_{i}_value"] = {0: "batch_size"} + return result + + @property + def outputs(self) -> dict[str, dict[int, str]]: # noqa: D102 + result: dict[str, dict[int, str]] = {"logits": {0: "batch_size"}} + for i in range(self._normalized_config.num_layers): + result[f"present_{i}_key"] = {0: "batch_size"} + result[f"present_{i}_value"] = {0: "batch_size"} + return result + + +class Florence2DecoderWrapper(WinMLDecoderWrapper): + """Export Florence-2's BART decoder and language-model head.""" + + _HF_MODEL_CLS = _NativeFlorence2ForConditionalGeneration + _IO_CONFIG_CLS = Florence2DecoderIOConfig + + def _invoke_hf(self, cache: Any, inputs: dict[str, torch.Tensor]) -> torch.Tensor: + cache_position = inputs["cache_position"].squeeze() + legacy_cache = tuple( + ( + layer.keys[:, :, :cache_position, :], + layer.values[:, :, :cache_position, :], + ) + for layer in cache.layers + ) + get_decoder = cast("Callable[[], _Florence2Decoder]", self.model.get_decoder) + decoder = get_decoder() + decoder_outputs = decoder( + input_ids=inputs["decoder_input_ids"], + attention_mask=inputs["decoder_attention_mask"][ + :, : cache_position + inputs["decoder_input_ids"].size(1) + ], + encoder_hidden_states=inputs["encoder_hidden_states"], + encoder_attention_mask=None, + past_key_values=legacy_cache, + use_cache=True, + return_dict=True, + ) + for index, (key, value, *_) in enumerate(decoder_outputs.past_key_values): + cache.captured[index] = ( + key[:, :, cache_position:, :], + value[:, :, cache_position:, :], + ) + return cast("_Florence2LanguageModel", self.model.language_model).lm_head( + decoder_outputs.last_hidden_state + ) + + +@register_composite_model("florence2", "image-to-text") +class WinMLFlorence2ImageToText(WinMLEncoderDecoderModel): + """Florence-2 image-to-text inference model.""" + + main_input_name = "pixel_values" + pipeline_capabilities = frozenset({PipelineCapability.COMBINED_IMAGE_TEXT_PROCESSOR}) + _SUB_MODEL_CONFIG: ClassVar[dict[str, str]] = { + "encoder": "image-feature-extraction", + "decoder": "text2text-generation", + } + + def __init__(self, sub_models: dict[str, Any], config: PretrainedConfig) -> None: + super().__init__(sub_models, config) + self.config.is_encoder_decoder = True + + def create_combined_processor(self, model_id: str) -> Any: + """Load the processor required by the declared preprocessing capability.""" + return _load_native_combined_processor( + model_id, + trust_remote_code=self._trust_remote_code, + ) + + @classmethod + def get_cache_class(cls) -> type: # noqa: D102 + return WinMLStaticCache + + @property + def generation_config(self) -> GenerationConfig: # noqa: D102 + if not hasattr(self, "_generation_config"): + from transformers import GenerationConfig + + text_config = self.config.text_config + self._generation_config = GenerationConfig( + decoder_start_token_id=text_config.decoder_start_token_id, + bos_token_id=text_config.bos_token_id, + eos_token_id=text_config.eos_token_id, + pad_token_id=text_config.pad_token_id, + max_new_tokens=self._max_dec - 1, + num_beams=1, + do_sample=False, + ) + return self._generation_config + + @generation_config.setter + def generation_config(self, value: Any) -> None: + self._generation_config = value + + +MODEL_CLASS_MAPPING: dict[tuple[str, str | None], type] = { + ("florence2", None): WinMLFlorence2ImageToText, + ("florence2", "image-to-text"): WinMLFlorence2ImageToText, + ("florence2", "feature-extraction"): Florence2EncoderWrapper, + ("florence2", "text2text-generation"): Florence2DecoderWrapper, +} + + +__all__ = [ + "FLORENCE2_CONFIG", + "MODEL_CLASS_MAPPING", + "Florence2DecoderIOConfig", + "Florence2DecoderWrapper", + "Florence2EncoderIOConfig", + "Florence2EncoderWrapper", + "WinMLFlorence2ImageToText", +] diff --git a/src/winml/modelkit/models/winml/__init__.py b/src/winml/modelkit/models/winml/__init__.py index c201275e8..c61e1e09e 100644 --- a/src/winml/modelkit/models/winml/__init__.py +++ b/src/winml/modelkit/models/winml/__init__.py @@ -181,6 +181,7 @@ def register_specialization(model_type: str, task: str, class_name: str) -> None from .base import WinMLModelForGenericTask, WinMLPreTrainedModel from .composite_model import ( COMPOSITE_MODEL_REGISTRY, + PipelineCapability, WinMLCompositeModel, register_composite_model, ) @@ -209,6 +210,7 @@ def register_specialization(model_type: str, task: str, class_name: str) -> None "TASK_TO_WINML_CLASS", "WINML_MODEL_CLASS_MAPPING", "ImageSegmentationOutput", + "PipelineCapability", "WinMLCache", "WinMLCompositeModel", "WinMLDecoderOnlyModel", diff --git a/src/winml/modelkit/models/winml/composite_model.py b/src/winml/modelkit/models/winml/composite_model.py index cc5d8e517..b20aa9610 100644 --- a/src/winml/modelkit/models/winml/composite_model.py +++ b/src/winml/modelkit/models/winml/composite_model.py @@ -41,6 +41,7 @@ from __future__ import annotations import logging +from enum import StrEnum from typing import TYPE_CHECKING, Any, ClassVar, cast import torch @@ -57,6 +58,12 @@ logger = logging.getLogger(__name__) +class PipelineCapability(StrEnum): + """Optional preprocessing contracts required by a composite model.""" + + COMBINED_IMAGE_TEXT_PROCESSOR = "combined-image-text-processor" + + # ========================================================================= # composite model Registry # ========================================================================= @@ -105,6 +112,7 @@ class WinMLCompositeModel(PreTrainedModel): """ _SUB_MODEL_CONFIG: ClassVar[dict[str, str]] = {} + pipeline_capabilities: ClassVar[frozenset[PipelineCapability]] = frozenset() def __init__( self, @@ -115,6 +123,7 @@ def __init__( self.sub_models = sub_models self.config = config self._device = device + self._trust_remote_code = False @classmethod def from_pretrained( @@ -201,7 +210,9 @@ def from_pretrained( **merged, ) - return cls(sub_models=sub_models, config=hf_config) + model = cls(sub_models=sub_models, config=hf_config) + model._trust_remote_code = trust_remote_code + return model @classmethod def from_onnx( @@ -213,6 +224,7 @@ def from_onnx( task: str | None = None, hf_config: PretrainedConfig | None = None, sub_model_kwargs: dict[str, dict[str, Any]] | None = None, + trust_remote_code: bool = False, **kwargs: Any, ) -> WinMLCompositeModel: """Load composite model from pre-built ONNX files. @@ -269,7 +281,9 @@ def from_onnx( if hf_config is None: raise ValueError("Composite model construction requires an HF config (hf_config).") - return resolved_cls(sub_models=sub_models, config=hf_config) + model = resolved_cls(sub_models=sub_models, config=hf_config) + model._trust_remote_code = trust_remote_code + return model @property def device(self) -> torch.device: diff --git a/src/winml/modelkit/models/winml/encoder_decoder.py b/src/winml/modelkit/models/winml/encoder_decoder.py index d488ae842..2bc0b19e1 100644 --- a/src/winml/modelkit/models/winml/encoder_decoder.py +++ b/src/winml/modelkit/models/winml/encoder_decoder.py @@ -188,11 +188,20 @@ def __init__( # Build {name: shape} lookups from ONNX I/O metadata enc_io = raw_encoder.io_config + enc_input_names = enc_io.get("input_names", []) + if not enc_input_names: + raise KeyError("Encoder ONNX I/O metadata is missing input_names.") enc_expected = dict( - zip(enc_io.get("input_names", []), enc_io.get("input_shapes", []), strict=False) + zip(enc_input_names, enc_io.get("input_shapes", []), strict=False) ) + self._enc_expected = enc_expected # Wrap encoder with auto-padding so all callsites just use self._encoder(...) - self._encoder = self._EncoderWithInputPadding(raw_encoder, enc_expected) + self._encoder = self._EncoderWithInputPadding( + raw_encoder, + enc_expected, + main_input_name=enc_input_names[0], + ) + self.encoder = self._encoder dec_io = self._decoder.io_config self._dec_expected = dict( @@ -230,10 +239,17 @@ class _EncoderWithInputPadding(torch.nn.Module): ``get_encoder()`` (GenerationMixin contract). """ - def __init__(self, encoder: Any, expected: dict[str, list[int]]) -> None: + def __init__( + self, + encoder: Any, + expected: dict[str, list[int]], + *, + main_input_name: str, + ) -> None: super().__init__() self._encoder = encoder self._expected = expected + self.main_input_name = main_input_name def forward(self, **kwargs: Any) -> BaseModelOutput: feeds = pad_inputs(kwargs, self._expected) @@ -247,6 +263,12 @@ def get_encoder(self) -> torch.nn.Module: def can_generate(self) -> bool: # noqa: D102 return True + def _validate_model_kwargs(self, model_kwargs: dict[str, Any]) -> None: + """Allow declared encoder ONNX inputs through GenerationMixin validation.""" + super()._validate_model_kwargs( + {name: value for name, value in model_kwargs.items() if name not in self._enc_expected} + ) + def prepare_inputs_for_generation( # type: ignore[override] # GenerationMixin's base signature differs; static-cache flow self, input_ids: torch.LongTensor, @@ -258,7 +280,8 @@ def prepare_inputs_for_generation( # type: ignore[override] # GenerationMixin' """Build decoder inputs for each generate() step.""" from .kv_cache import WinMLCache - if isinstance(past_key_values, WinMLCache) and past_key_values.get_seq_length() > 0: + cache = getattr(past_key_values, "self_attention_cache", past_key_values) + if isinstance(cache, WinMLCache) and cache.step > 0: decoder_input_ids = input_ids[:, -1:] else: decoder_input_ids = input_ids diff --git a/src/winml/modelkit/onnx/__init__.py b/src/winml/modelkit/onnx/__init__.py index 762603498..9f09ca665 100644 --- a/src/winml/modelkit/onnx/__init__.py +++ b/src/winml/modelkit/onnx/__init__.py @@ -27,7 +27,7 @@ ) from .metadata import capture_metadata, restore_metadata from .persistence import ONNXSaveError, cleanup_onnx, load_onnx, save_onnx -from .shape import infer_onnx_shapes, infer_shapes +from .shape import infer_onnx_shapes, infer_shapes, infer_symbolic_shapes from .utils import EXTERNAL_DATA_THRESHOLD, check_onnx_model, get_model_size, strip_node_attrs @@ -50,6 +50,7 @@ "get_onnx_model_hash", "infer_onnx_shapes", "infer_shapes", + "infer_symbolic_shapes", "is_compiled_onnx", "is_quantized_onnx", "load_onnx", diff --git a/src/winml/modelkit/onnx/io.py b/src/winml/modelkit/onnx/io.py index 1c58362cf..62f742ea2 100644 --- a/src/winml/modelkit/onnx/io.py +++ b/src/winml/modelkit/onnx/io.py @@ -21,6 +21,7 @@ import json import logging +import math from dataclasses import dataclass from pathlib import Path from typing import Any @@ -31,6 +32,8 @@ logger = logging.getLogger(__name__) ShapeDim = int | str +DummyValue = int | float +DummyValueRun = tuple[int, DummyValue] # ============================================================================= @@ -53,6 +56,9 @@ class InputTensorSpec: Populated by resolve_export_config() via Optimum's interceptor. Integer semantics: torch.randint(min, max) — max is exclusive. Float semantics: uniform in [min, max). + dummy_value_runs: Optional run-length encoded values for deterministic + semantic dummy input generation. Each item is ``(count, value)``. + The runs must fill the concrete tensor shape exactly. Example: # Vision model input @@ -69,6 +75,28 @@ class InputTensorSpec: dtype: str | None = None # "float32", "float16", "int64", "int32", etc. shape: tuple[ShapeDim, ...] | None = None value_range: tuple[float, float] | None = None # (min, max_exclusive) + dummy_value_runs: tuple[DummyValueRun, ...] | None = None + + def __post_init__(self) -> None: + """Normalize and validate serialized semantic dummy values.""" + if self.dummy_value_runs is None: + return + + normalized_runs: list[DummyValueRun] = [] + for run in self.dummy_value_runs: + if not isinstance(run, (tuple, list)) or len(run) != 2: + raise TypeError("dummy_value_runs entries must be (count, value) pairs") + count, value = run + if not isinstance(count, int) or isinstance(count, bool): + raise TypeError("dummy_value_runs counts must be integers") + if count <= 0: + raise ValueError("dummy_value_runs counts must be positive") + if not isinstance(value, (int, float)) or isinstance(value, bool): + raise TypeError("dummy_value_runs values must be numeric") + if isinstance(value, float) and not math.isfinite(value): + raise ValueError("dummy_value_runs values must be finite") + normalized_runs.append((count, value)) + self.dummy_value_runs = tuple(normalized_runs) def to_tensor(self) -> Any: """Generate a dummy tensor from this spec. @@ -102,6 +130,26 @@ def to_tensor(self) -> Any: concrete_shape = self.concrete_shape() + if self.dummy_value_runs is not None: + if sum(count for count, _ in self.dummy_value_runs) != math.prod(concrete_shape): + raise ValueError( + f"dummy_value_runs for '{self.name}' must fill the concrete tensor shape" + ) + if not torch_dtype.is_floating_point and any( + not isinstance(value, int) for _, value in self.dummy_value_runs + ): + raise TypeError( + f"dummy_value_runs for integer tensor '{self.name}' must contain integers" + ) + + tensor = torch.empty(concrete_shape, dtype=torch_dtype) + flat_tensor = tensor.reshape(-1) + start = 0 + for count, value in self.dummy_value_runs: + flat_tensor[start : start + count] = value + start += count + return tensor + if self.value_range is not None: lo, hi = self.value_range if torch_dtype.is_floating_point: @@ -133,6 +181,35 @@ def concrete_shape(self) -> tuple[int, ...]: ) return tuple(concrete) + @staticmethod + def compact_dummy_value_runs(tensor: Any) -> tuple[DummyValueRun, ...] | None: + """Return compact run-length encoded values for a generated dummy tensor. + + Random tensors are intentionally omitted: their run count equals their + element count, so serializing them would bloat recipes without preserving + useful semantics. + """ + import torch + + if not isinstance(tensor, torch.Tensor) or tensor.numel() < 2: + return None + if tensor.dtype == torch.bool or tensor.is_complex(): + return None + + flat_tensor = tensor.detach().cpu().reshape(-1) + values, counts = torch.unique_consecutive(flat_tensor, return_counts=True) + if values.numel() >= flat_tensor.numel(): + return None + + runs: list[DummyValueRun] = [] + for count, value in zip(counts.tolist(), values.tolist(), strict=True): + if not isinstance(value, (int, float)) or isinstance(value, bool): + return None + if isinstance(value, float) and not math.isfinite(value): + return None + runs.append((count, value)) + return tuple(runs) + def to_dict(self) -> dict[str, Any]: """Convert to dictionary, excluding None values.""" result: dict[str, Any] = {} @@ -144,6 +221,8 @@ def to_dict(self) -> dict[str, Any]: result["shape"] = self.shape if self.value_range is not None: result["value_range"] = list(self.value_range) + if self.dummy_value_runs is not None: + result["dummy_value_runs"] = [list(run) for run in self.dummy_value_runs] return result @classmethod @@ -159,6 +238,7 @@ def from_dict(cls, data: dict[str, Any]) -> InputTensorSpec: dtype=data.get("dtype"), shape=shape, value_range=value_range, + dummy_value_runs=data.get("dummy_value_runs"), ) diff --git a/src/winml/modelkit/onnx/shape.py b/src/winml/modelkit/onnx/shape.py index 53f6925ca..e9aa5835b 100644 --- a/src/winml/modelkit/onnx/shape.py +++ b/src/winml/modelkit/onnx/shape.py @@ -17,9 +17,13 @@ from __future__ import annotations import logging +import multiprocessing as mp +import os +import queue import tempfile +import traceback from pathlib import Path -from typing import cast +from typing import Any, cast import onnx @@ -68,18 +72,7 @@ def _run_inference(model: onnx.ModelProto) -> onnx.ModelProto: """ # Try symbolic first (handles com.microsoft ops from ORT fusion/quantization) try: - from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference - - return cast( - "onnx.ModelProto", - SymbolicShapeInference.infer_shapes( - model, - int_max=2**31 - 1, - auto_merge=False, - guess_output_rank=False, - verbose=0, - ), - ) + return infer_symbolic_shapes(model) except Exception as e: logger.debug("Symbolic shape inference failed: %s", e) @@ -95,6 +88,63 @@ def _run_inference(model: onnx.ModelProto) -> onnx.ModelProto: return model +def _infer_symbolic_shapes_worker( + model_bytes: bytes, + scratch_dir: str, + result_queue: Any, +) -> None: + """Run symbolic inference in a worker's isolated current directory.""" + from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference + + try: + os.chdir(scratch_dir) + result = cast( + "onnx.ModelProto", + SymbolicShapeInference.infer_shapes( + onnx.load_model_from_string(model_bytes), + int_max=2**31 - 1, + auto_merge=False, + guess_output_rank=False, + verbose=0, + ), + ) + result_queue.put((True, result.SerializeToString())) + except Exception: + result_queue.put((False, traceback.format_exc())) + + +def infer_symbolic_shapes(model: onnx.ModelProto) -> onnx.ModelProto: + """Run ORT symbolic inference in an isolated worker process.""" + with tempfile.TemporaryDirectory(prefix="winmlcli_shape_") as scratch_dir: + context = mp.get_context("spawn") + result_queue = context.Queue() + worker = context.Process( + target=_infer_symbolic_shapes_worker, + args=(model.SerializeToString(), scratch_dir, result_queue), + ) + worker.start() + try: + while True: + try: + success, payload = result_queue.get(timeout=0.1) + break + except queue.Empty: + if not worker.is_alive(): + raise RuntimeError( + "Symbolic shape inference worker " + f"exited with code {worker.exitcode} without a result" + ) from None + worker.join() + finally: + result_queue.close() + result_queue.join_thread() + worker.close() + + if not success: + raise RuntimeError(f"Symbolic shape inference worker failed:\n{payload}") + return onnx.load_model_from_string(payload) + + def infer_onnx_shapes( model: onnx.ModelProto, check_type: bool = False, diff --git a/src/winml/modelkit/utils/config_utils.py b/src/winml/modelkit/utils/config_utils.py index 26e514fe0..6e86b39d8 100644 --- a/src/winml/modelkit/utils/config_utils.py +++ b/src/winml/modelkit/utils/config_utils.py @@ -100,6 +100,9 @@ def _merge_dataclass(base: T, overrides: dict[str, Any]) -> T: current[key] = None elif current_val is None: # Base is None, override has value - use override + if isinstance(value, list): + current[key] = _reconstruct_list_items(base, key, value) + continue # Try to construct from dict if nested config field_type = _get_field_type(base, key) if field_type and isinstance(value, dict): @@ -118,7 +121,7 @@ def _merge_dataclass(base: T, overrides: dict[str, Any]) -> T: current[key] = _merge_into(current_val, value) elif isinstance(value, list) and isinstance(current_val, list): # List - replace entirely (no merge) - current[key] = value + current[key] = _reconstruct_list_items(base, key, value) else: # Primitive - override current[key] = value @@ -160,3 +163,39 @@ def _get_field_type(obj: Any, field_name: str) -> type | None: if arg is not type(None): return cast("type", arg) return resolved if isinstance(resolved, type) else None + + +def _reconstruct_list_items(obj: Any, field_name: str, values: list[Any]) -> list[Any]: + """Rebuild typed list items from serialized configuration data.""" + try: + annotation = typing.get_type_hints(type(obj)).get(field_name) + except (NameError, AttributeError): + annotation = None + + annotation_args = typing.get_args(annotation) + if type(None) in annotation_args: + annotation = next( + (arg for arg in annotation_args if arg is not type(None)), + None, + ) + + if typing.get_origin(annotation) is not list: + return values + + item_types = typing.get_args(annotation) + if len(item_types) != 1 or not isinstance(item_types[0], type): + return values + + item_type = item_types[0] + if not hasattr(item_type, "from_dict") and not dataclasses.is_dataclass(item_type): + return values + + reconstructed: list[Any] = [] + for value in values: + if not isinstance(value, dict): + reconstructed.append(value) + elif hasattr(item_type, "from_dict"): + reconstructed.append(item_type.from_dict(value)) + else: + reconstructed.append(item_type(**value)) + return reconstructed diff --git a/tests/integration/export/test_florence2_processor_contract.py b/tests/integration/export/test_florence2_processor_contract.py new file mode 100644 index 000000000..9df902631 --- /dev/null +++ b/tests/integration/export/test_florence2_processor_contract.py @@ -0,0 +1,237 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Native Florence-2 processor contract checks for the exported encoder.""" + +from __future__ import annotations + +import gc +from pathlib import Path +from shutil import rmtree +from typing import TYPE_CHECKING, NamedTuple + +import numpy as np +import onnx +import onnxruntime as ort +import pytest +import torch +from click.testing import CliRunner +from PIL import Image +from transformers import AutoConfig + +from winml.modelkit.commands.build import build +from winml.modelkit.eval.metrics.tensor_similarity import TensorSimilarityMetric +from winml.modelkit.export import resolve_io_specs +from winml.modelkit.models.hf.florence2 import ( + _load_native_combined_processor, + _NativeFlorence2ForConditionalGeneration, +) + + +if TYPE_CHECKING: + from collections.abc import Iterator + + +_MODEL_ID = "microsoft/Florence-2-base" +_RECIPE_DIR = Path("examples/recipes/microsoft_Florence-2-base") +_ARTIFACT_ROOT = Path("temp/microsoft_Florence-2-base/processor-contract") + +pytestmark = [pytest.mark.integration, pytest.mark.network, pytest.mark.slow] + + +class FlorenceArtifacts(NamedTuple): + """Fresh model and component artifacts owned by the module fixture.""" + + model_dir: Path + encoder_onnx_path: Path + encoder_export_path: Path + decoder_onnx_path: Path + + +def _build_recipe_component(recipe_path: Path, model_dir: Path, output_dir: Path) -> Path: + """Build one recipe component through the public precision-variant command path.""" + result = CliRunner().invoke( + build, + [ + "--config", + str(recipe_path), + "--model", + str(model_dir), + "--output-dir", + str(output_dir), + "--precision", + "fp16", + "--no-compile", + "--rebuild", + ], + obj={"debug": False}, + ) + assert result.exit_code == 0, result.output + onnx_path = output_dir / "model.onnx" + assert onnx_path.is_file() + return onnx_path + + +@pytest.fixture(scope="module") +def florence_artifacts() -> Iterator[FlorenceArtifacts]: + """Download and build component artifacts once in fixture-owned storage.""" + from huggingface_hub import snapshot_download + + root = _ARTIFACT_ROOT + if root.exists(): + rmtree(root) + root.mkdir(parents=True) + model_dir = root / "model" + build_dir = root / "build" + try: + snapshot_download( + _MODEL_ID, + local_dir=model_dir, + cache_dir=root / "hf-cache", + ) + + encoder_output_dir = build_dir / "encoder" + encoder_onnx_path = _build_recipe_component( + _RECIPE_DIR / "image-to-text_fp16_config_encoder.json", + model_dir, + encoder_output_dir, + ) + decoder_onnx_path = _build_recipe_component( + _RECIPE_DIR / "image-to-text_fp16_config_decoder.json", + model_dir, + build_dir / "decoder", + ) + yield FlorenceArtifacts( + model_dir=model_dir, + encoder_onnx_path=encoder_onnx_path, + encoder_export_path=encoder_output_dir / "export.onnx", + decoder_onnx_path=decoder_onnx_path, + ) + finally: + gc.collect() + rmtree(root) + if not any(root.parent.iterdir()): + root.parent.rmdir() + + +def _session_inputs(session: ort.InferenceSession, batch: dict[str, torch.Tensor]) -> dict: + input_names = {input_.name for input_ in session.get_inputs()} + return { + name: value.detach().cpu().numpy() for name, value in batch.items() if name in input_names + } + + +def test_native_model_loads_the_complete_checkpoint(florence_artifacts: FlorenceArtifacts) -> None: + """Independent native loads must not retain random initialized parameters.""" + first = _NativeFlorence2ForConditionalGeneration.from_pretrained( + florence_artifacts.model_dir, + trust_remote_code=True, + ) + second = _NativeFlorence2ForConditionalGeneration.from_pretrained( + florence_artifacts.model_dir, + trust_remote_code=True, + ) + + first_state = first.state_dict() + second_state = second.state_dict() + assert first_state.keys() == second_state.keys() + for key in first_state: + assert torch.equal(first_state[key], second_state[key]), key + + +def test_native_processor_matches_exported_encoder_contract( + florence_artifacts: FlorenceArtifacts, +) -> None: + config = AutoConfig.from_pretrained(florence_artifacts.model_dir) + processor = _load_native_combined_processor( + str(florence_artifacts.model_dir), + trust_remote_code=True, + ) + batch = processor( + text="", + images=Image.new("RGB", (768, 768)), + return_tensors="pt", + ) + specs = resolve_io_specs("florence2", "feature-extraction", config) + shapes = dict(zip(specs["input_names"], specs["input_shapes"], strict=True)) + + assert set(batch) == set(shapes) + for name, shape in shapes.items(): + assert tuple(batch[name].shape) == shape + assert tuple(batch["input_ids"].shape) == (1, 8) + + +def test_native_padded_prompts_match_exported_encoder( + florence_artifacts: FlorenceArtifacts, +) -> None: + processor = _load_native_combined_processor( + str(florence_artifacts.model_dir), + trust_remote_code=True, + ) + batch = processor( + text=["", ""], + images=[Image.new("RGB", (768, 768)), Image.new("RGB", (768, 768))], + padding=True, + return_tensors="pt", + ) + assert torch.any(batch["attention_mask"] == 0) + session = ort.InferenceSession( + str(florence_artifacts.encoder_export_path), providers=["CPUExecutionProvider"] + ) + assert {input_.name for input_ in session.get_inputs()} == set(batch) + exported = session.run(None, _session_inputs(session, batch))[0] + optimized_session = ort.InferenceSession( + str(florence_artifacts.encoder_onnx_path), providers=["CPUExecutionProvider"] + ) + optimized = optimized_session.run(None, _session_inputs(optimized_session, batch))[0] + + optimized_model = onnx.load(florence_artifacts.encoder_onnx_path, load_external_data=False) + assert any( + initializer.data_type == onnx.TensorProto.FLOAT16 + for initializer in optimized_model.graph.initializer + ) + assert optimized.shape == exported.shape + assert optimized.dtype == exported.dtype + assert np.isfinite(optimized).all() + + raw_metrics = TensorSimilarityMetric() + raw_metrics.update(optimized, exported) + metrics = raw_metrics.compute() + assert {"cosine_similarity_mean", "max_abs_diff_mean"} <= metrics.keys() + + +def test_decoder_accepts_real_caption_encoder_states(florence_artifacts: FlorenceArtifacts) -> None: + processor = _load_native_combined_processor( + str(florence_artifacts.model_dir), + trust_remote_code=True, + ) + batch = processor( + text="", + images=Image.new("RGB", (768, 768)), + return_tensors="pt", + ) + encoder_session = ort.InferenceSession( + str(florence_artifacts.encoder_onnx_path), providers=["CPUExecutionProvider"] + ) + encoder_hidden_states = encoder_session.run( + None, _session_inputs(encoder_session, batch) + )[0] + decoder_session = ort.InferenceSession( + str(florence_artifacts.decoder_onnx_path), providers=["CPUExecutionProvider"] + ) + decoder_inputs = { + "decoder_input_ids": torch.zeros((1, 1), dtype=torch.int32).numpy(), + "encoder_hidden_states": encoder_hidden_states, + "decoder_attention_mask": torch.ones((1, 1024), dtype=torch.long).numpy(), + "cache_position": torch.zeros((1,), dtype=torch.long).numpy(), + } + for index in range(6): + decoder_inputs[f"past_{index}_key"] = torch.zeros((1, 12, 1024, 64)).numpy() + decoder_inputs[f"past_{index}_value"] = torch.zeros((1, 12, 1024, 64)).numpy() + + outputs = decoder_session.run(None, decoder_inputs) + + assert encoder_hidden_states.shape == (1, 585, 768) + assert outputs[0].shape[:2] == (1, 1) diff --git a/tests/unit/commands/test_config_composite_resolution.py b/tests/unit/commands/test_config_composite_resolution.py index a364af2d6..7ef36f46c 100644 --- a/tests/unit/commands/test_config_composite_resolution.py +++ b/tests/unit/commands/test_config_composite_resolution.py @@ -23,7 +23,7 @@ from unittest.mock import patch import pytest -from transformers import BartConfig, Qwen3Config, T5Config +from transformers import BartConfig, Florence2Config, Qwen3Config, T5Config from winml.modelkit.commands.config import ( _resolve_composite_model_components as _resolve, @@ -31,6 +31,9 @@ from winml.modelkit.loader.resolution import ( _composite_components_for_task as _serve, ) +from winml.modelkit.loader.resolution import ( + resolve_task, +) # ============================================================================= @@ -57,6 +60,20 @@ def test_blip_image_to_text_expands_to_composite() -> None: assert "encoder" in components and "decoder" in components +def test_florence2_default_resolution_matches_image_to_text_composite() -> None: + config = Florence2Config() + config.architectures = ["Florence2ForConditionalGeneration"] + + resolution = resolve_task(config) + + assert resolution.task == "image-to-text" + assert resolution.composite == { + "encoder": "image-feature-extraction", + "decoder": "text2text-generation", + } + assert _resolve(None, "florence2", None) == resolution.composite + + def test_bart_text_classification_stays_single() -> None: """facebook/bart-large-mnli (BartForSequenceClassification) detects text-classification; it must NOT expand to a seq2seq composite -- consistent diff --git a/tests/unit/commands/test_inspect_cli.py b/tests/unit/commands/test_inspect_cli.py index be88be42c..735e4367b 100644 --- a/tests/unit/commands/test_inspect_cli.py +++ b/tests/unit/commands/test_inspect_cli.py @@ -91,6 +91,7 @@ def test_help_shows_all_options(self, runner: CliRunner) -> None: "-m", "--format", "-f", + "--trust-remote-code", "--verbose", "-v", "--task", @@ -235,6 +236,32 @@ def test_task_override_passed_to_api( _, call_kwargs = mock_api.call_args assert call_kwargs["task_override"] == "fill-mask" + @pytest.mark.parametrize( + ("args", "expected_trust_remote_code"), + [ + ([], False), + (["--trust-remote-code"], True), + ], + ) + def test_trust_remote_code_passed_to_api( + self, + runner: CliRunner, + mock_inspect_result: MagicMock, + args: list[str], + expected_trust_remote_code: bool, + ) -> None: + """The CLI must pass its explicit or default trust consent to the API.""" + from winml.modelkit.commands.inspect import inspect + + with ( + patch(_INSPECT_MODEL, return_value=mock_inspect_result) as mock_api, + patch(_OUTPUT_TABLE), + ): + result = runner.invoke(inspect, ["-m", "test", *args], obj={}) + + assert result.exit_code == 0, result.output + assert mock_api.call_args.kwargs["trust_remote_code"] is expected_trust_remote_code + def test_hierarchy_flag_passed_to_api( self, runner: CliRunner, diff --git a/tests/unit/eval/test_image_to_text_evaluator.py b/tests/unit/eval/test_image_to_text_evaluator.py index f20a3665b..9e66d1165 100644 --- a/tests/unit/eval/test_image_to_text_evaluator.py +++ b/tests/unit/eval/test_image_to_text_evaluator.py @@ -12,7 +12,7 @@ from winml.modelkit.eval.image_to_text_evaluator import WinMLImageToTextEvaluator -def make_evaluator(columns_mapping=None): +def make_evaluator(columns_mapping=None, prompt=None): """Instantiate evaluator with mocked dataset + pipeline.""" from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig @@ -34,15 +34,30 @@ def make_evaluator(columns_mapping=None): config = WinMLEvaluationConfig( model_id="microsoft/trocr-base-handwritten", task="image-to-text", + prompt=prompt, dataset=DatasetConfig(path="Teklia/IAM-line", columns_mapping=mapping), ) - with patch("datasets.load_dataset", return_value=mock_ds), \ - patch("transformers.pipeline", return_value=mock_pipe): + with ( + patch("datasets.load_dataset", return_value=mock_ds), + patch( + "winml.modelkit.eval.base_evaluator.WinMLEvaluator.prepare_pipeline", + return_value=mock_pipe, + ), + ): return WinMLImageToTextEvaluator(config, model) class TestInit: + def test_config_serializes_prompt(self): + """Evaluation configs preserve task prompts for prompt-aware pipelines.""" + from winml.modelkit.eval import WinMLEvaluationConfig + + config = WinMLEvaluationConfig(task="image-to-text", prompt="") + + assert config.to_dict()["prompt"] == "" + assert WinMLEvaluationConfig.from_dict(config.to_dict()).prompt == "" + def test_default_columns(self): ev = make_evaluator() assert ev._image_col == "image" @@ -99,6 +114,16 @@ def test_perfect_predictions(self): assert result["n_samples"] == 2 assert "cider" in result + def test_passes_configured_prompt_to_pipeline(self): + """Prompt-aware image-to-text pipelines receive the configured prompt.""" + ev = make_evaluator(prompt="") + ev.data = [{"image": "img1", "text": "caption"}] + ev.pipe = MagicMock(return_value=[{"generated_text": "caption"}]) + + ev.compute() + + ev.pipe.assert_called_once_with("img1", prompt="") + def test_dict_output_shape(self): """Pipeline may also return a single dict (not a list).""" ev = make_evaluator() diff --git a/tests/unit/eval/test_pipeline_factory.py b/tests/unit/eval/test_pipeline_factory.py new file mode 100644 index 000000000..b51b87ad5 --- /dev/null +++ b/tests/unit/eval/test_pipeline_factory.py @@ -0,0 +1,56 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Tests that evaluators use the shared capability-aware pipeline factory.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from winml.modelkit.eval import WinMLEvaluationConfig +from winml.modelkit.eval.base_evaluator import WinMLEvaluator +from winml.modelkit.inference.pipeline import PipelineCapability + + +def test_evaluator_delegates_capability_aware_pipeline_creation(monkeypatch) -> None: + model = SimpleNamespace( + pipeline_capabilities=frozenset( + {PipelineCapability.COMBINED_IMAGE_TEXT_PROCESSOR} + ) + ) + config = WinMLEvaluationConfig(model_id="local-model", task="image-to-text") + evaluator = object.__new__(WinMLEvaluator) + evaluator.model = model + evaluator.config = config + calls: list[tuple[str, Any, str | None]] = [] + expected_pipeline = object() + + def create_pipeline(task: str, pipeline_model: Any, model_id: str | None) -> object: + calls.append((task, pipeline_model, model_id)) + return expected_pipeline + + monkeypatch.setattr("winml.modelkit.inference.pipeline.create_pipeline", create_pipeline) + + assert evaluator.prepare_pipeline() is expected_pipeline + assert calls == [("image-to-text", model, "local-model")] + + +def test_pipeline_class_ignores_synthesized_mock_capabilities() -> None: + """Models without a declared capability contract use the default pipeline.""" + from winml.modelkit.inference.pipeline import _pipeline_class_for + + assert _pipeline_class_for(MagicMock()) is None + + +def test_pipeline_class_rejects_declared_invalid_capabilities() -> None: + """An explicitly declared contract must remain strictly validated.""" + from winml.modelkit.inference.pipeline import _pipeline_class_for + + with pytest.raises(TypeError, match="must be a frozenset"): + _pipeline_class_for(SimpleNamespace(pipeline_capabilities=set())) # type: ignore[arg-type] diff --git a/tests/unit/eval/test_run_eval_script.py b/tests/unit/eval/test_run_eval_script.py index 900a61109..e0c4641d8 100644 --- a/tests/unit/eval/test_run_eval_script.py +++ b/tests/unit/eval/test_run_eval_script.py @@ -18,6 +18,7 @@ import json import sys from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest @@ -836,6 +837,28 @@ def test_composite_builds_each_role(self, run_eval, tmp_path): assert len(captured) == 2 # one winml build per component assert all("--use-cache" in call for call in captured) + def test_composite_build_forwards_variant_precision_to_each_component( + self, run_eval + ): + variant = SimpleNamespace( + precision="fp16", + components=[ + SimpleNamespace(path=Path("encoder.json"), role="encoder"), + SimpleNamespace(path=Path("decoder.json"), role="decoder"), + ], + ) + captured: list[list[str]] = [] + with ( + patch.object(run_eval, "_run_subprocess", side_effect=self._fake_subprocess(captured)), + patch.object(run_eval, "_extract_onnx_path", side_effect=lambda *a: "m.onnx"), + ): + run_eval._run_recipe_build(_entry(), variant, 300, Path("temp") / "out") + + assert len(captured) == 2 + assert all( + call[call.index("--precision") + 1] == variant.precision for call in captured + ) + def test_build_failure_reported(self, run_eval, tmp_path): variant = self._variant(run_eval, tmp_path, composite=False) diff --git a/tests/unit/export/test_config_validation.py b/tests/unit/export/test_config_validation.py index 47cdca59f..6bc565473 100644 --- a/tests/unit/export/test_config_validation.py +++ b/tests/unit/export/test_config_validation.py @@ -10,6 +10,7 @@ from __future__ import annotations +import json import logging import pytest @@ -19,6 +20,7 @@ OutputTensorSpec, WinMLExportConfig, ) +from winml.modelkit.utils.config_utils import merge_config # ============================================================================= @@ -242,6 +244,44 @@ def test_no_clean_no_warning(self, caplog): class TestInputTensorSpecRoundtrip: """InputTensorSpec serialization roundtrip.""" + def test_semantic_dummy_value_runs_survive_recipe_merge_and_export(self): + """Compact semantic values survive JSON recipe parsing and typed config reconstruction.""" + generated = WinMLExportConfig( + input_tensors=[ + InputTensorSpec( + name="tokens", + dtype="int64", + shape=(1, 5), + dummy_value_runs=((3, 17), (2, 0)), + ) + ] + ) + recipe = json.loads(json.dumps(generated.to_dict())) + + parsed_recipe = WinMLExportConfig.from_dict(recipe) + restored = merge_config(generated, parsed_recipe) + + assert restored.generate_dummy_inputs()["tokens"].tolist() == [[17, 17, 17, 0, 0]] + + def test_semantic_dummy_value_runs_reject_invalid_entries(self): + """Serialized runs must have positive integer counts and numeric values.""" + with pytest.raises(ValueError, match="counts must be positive"): + InputTensorSpec(dummy_value_runs=((0, 1),)) + with pytest.raises(TypeError, match="values must be numeric"): + InputTensorSpec(dummy_value_runs=((1, "not-a-number"),)) # type: ignore[arg-type] + + def test_semantic_dummy_value_runs_must_fill_tensor_shape(self): + """Semantic values cannot silently create a partially initialized tensor.""" + spec = InputTensorSpec( + name="tokens", + dtype="int64", + shape=(1, 3), + dummy_value_runs=((2, 1),), + ) + + with pytest.raises(ValueError, match="must fill the concrete tensor shape"): + spec.to_tensor() + def test_full_roundtrip(self): original = InputTensorSpec(name="pixel_values", dtype="float32", shape=(1, 3, 224, 224)) restored = InputTensorSpec.from_dict(original.to_dict()) diff --git a/tests/unit/export/test_florence2_onnx_config.py b/tests/unit/export/test_florence2_onnx_config.py new file mode 100644 index 000000000..2c46e4281 --- /dev/null +++ b/tests/unit/export/test_florence2_onnx_config.py @@ -0,0 +1,459 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Tests for Florence-2 split image-to-text export.""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace +from typing import ClassVar + +import torch +from optimum.exporters.tasks import TasksManager +from transformers.modeling_outputs import BaseModelOutput + +from winml.modelkit.config import WinMLBuildConfig +from winml.modelkit.export import generate_dummy_inputs, resolve_io_specs +from winml.modelkit.export.config import _resolve_export_config_from_specs +from winml.modelkit.models import HF_MODEL_CLASS_MAPPING + + +_RECIPE_DIR = Path("examples/recipes/microsoft_Florence-2-base") + + +def _florence2_config(): + from transformers import Florence2Config + + return Florence2Config() + + +def _florence2_base_config(): + from transformers import Florence2Config + + return Florence2Config( + text_config={ + "d_model": 768, + "decoder_attention_heads": 12, + "decoder_layers": 6, + "vocab_size": 51289, + } + ) + + +def _load_recipe(name: str) -> dict: + return json.loads((_RECIPE_DIR / name).read_text()) + + +class TestFlorence2Registration: + def test_class_mapping_contains_encoder(self) -> None: + assert ("florence2", "feature-extraction") in HF_MODEL_CLASS_MAPPING + + def test_class_mapping_contains_decoder(self) -> None: + assert ("florence2", "text2text-generation") in HF_MODEL_CLASS_MAPPING + + def test_encoder_config_registered(self) -> None: + constructor = TasksManager.get_exporter_config_constructor( + exporter="onnx", + model_type="florence2", + task="feature-extraction", + library_name="transformers", + ) + + assert constructor.func.__name__ == "Florence2EncoderIOConfig" + + def test_decoder_config_registered(self) -> None: + constructor = TasksManager.get_exporter_config_constructor( + exporter="onnx", + model_type="florence2", + task="text2text-generation", + library_name="transformers", + ) + + assert constructor.func.__name__ == "Florence2DecoderIOConfig" + + def test_composite_registered(self) -> None: + from winml.modelkit.models.winml.composite_model import COMPOSITE_MODEL_REGISTRY + + assert ("florence2", "image-to-text") in COMPOSITE_MODEL_REGISTRY + + +class TestFlorence2IO: + def test_encoder_accepts_image_and_prompt_inputs(self) -> None: + specs = resolve_io_specs("florence2", "feature-extraction", _florence2_config()) + + assert specs["input_names"] == ["input_ids", "pixel_values", "attention_mask"] + assert specs["output_names"] == ["last_hidden_state"] + assert specs["input_shapes"][0] == (1, 8) + assert specs["dummy_value_runs"]["input_ids"] == ((8, 0),) + + def test_decoder_has_static_cache_io(self) -> None: + specs = resolve_io_specs("florence2", "text2text-generation", _florence2_config()) + assert {"decoder_input_ids", "encoder_hidden_states", "cache_position"} <= set( + specs["input_names"] + ) + assert specs["output_names"][0] == "logits" + + def test_decoder_metadata_marks_encoder_source_axes_dynamic(self) -> None: + specs = resolve_io_specs("florence2", "text2text-generation", _florence2_config()) + + assert specs["dynamic_axes"]["encoder_hidden_states"] == { + 0: "batch_size", + 1: "sequence_length", + } + assert specs["dynamic_axes"]["decoder_input_ids"] == {0: "batch_size"} + assert specs["dynamic_axes"]["decoder_attention_mask"] == {0: "batch_size"} + assert specs["dynamic_axes"]["past_0_key"] == {0: "batch_size"} + + def test_encoder_export_config_recreates_generator_inputs(self) -> None: + """Semantic dummy values generated by an ONNX config survive export config construction.""" + config = _florence2_config() + generated_inputs = generate_dummy_inputs("florence2", "feature-extraction", config) + + export_config = _resolve_export_config_from_specs( + "florence2", + "feature-extraction", + config, + ) + + recreated_inputs = export_config.generate_dummy_inputs() + for name, generated in generated_inputs.items(): + assert torch.equal(recreated_inputs[name], generated) + + def test_encoder_prepends_image_features_to_text_embeddings(self) -> None: + """The upstream model combines its image embeddings with prompt embeddings.""" + config = SimpleNamespace() + + def embed(input_ids: torch.Tensor) -> torch.Tensor: + return torch.zeros((*input_ids.shape, 2)) + + def merge( + image_features: torch.Tensor, inputs_embeds: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + combined = torch.cat((image_features, inputs_embeds), dim=1) + return combined, torch.ones(combined.shape[:2], dtype=torch.long) + + fake_model = SimpleNamespace( + get_input_embeddings=lambda: embed, + _encode_image=lambda pixel_values: torch.zeros((1, 2, 2)), + _merge_input_ids_with_image_features=merge, + get_encoder=lambda: lambda **kwargs: SimpleNamespace( + last_hidden_state=kwargs["inputs_embeds"] + ), + ) + wrapper_class = HF_MODEL_CLASS_MAPPING[("florence2", "feature-extraction")] + wrapper = wrapper_class(fake_model, config) + + output = wrapper( + input_ids=torch.tensor([[5, 5, 0]]), + pixel_values=torch.zeros((1, 3, 2, 2)), + attention_mask=torch.ones((1, 3), dtype=torch.long), + ) + + assert output.shape == (1, 5, 2) + + def test_encoder_preserves_padded_text_masks_after_image_prefix(self) -> None: + """Image-token masks must be prefixed without unmasking padded prompt tokens.""" + config = SimpleNamespace() + received: dict[str, torch.Tensor] = {} + + def embed(input_ids: torch.Tensor) -> torch.Tensor: + return torch.zeros((*input_ids.shape, 2)) + + def merge( + image_features: torch.Tensor, inputs_embeds: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + combined = torch.cat((image_features, inputs_embeds), dim=1) + return combined, torch.ones(combined.shape[:2], dtype=torch.long) + + def encoder(**kwargs: torch.Tensor) -> SimpleNamespace: + received.update(kwargs) + return SimpleNamespace(last_hidden_state=kwargs["inputs_embeds"]) + + fake_model = SimpleNamespace( + get_input_embeddings=lambda: embed, + _encode_image=lambda pixel_values: torch.zeros((pixel_values.size(0), 2, 2)), + _merge_input_ids_with_image_features=merge, + get_encoder=lambda: encoder, + ) + wrapper_class = HF_MODEL_CLASS_MAPPING[("florence2", "feature-extraction")] + wrapper = wrapper_class(fake_model, config) + attention_mask = torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long) + + wrapper( + input_ids=torch.tensor([[5, 6, 0], [7, 0, 0]]), + pixel_values=torch.zeros((2, 3, 2, 2)), + attention_mask=attention_mask, + ) + + assert torch.equal( + received["attention_mask"], + torch.tensor([[1, 1, 1, 1, 0], [1, 1, 1, 0, 0]], dtype=torch.long), + ) + + +class TestFlorence2DecoderWrapper: + @staticmethod + def _cache(layer_count: int = 2) -> SimpleNamespace: + return SimpleNamespace( + layers=[ + SimpleNamespace( + keys=torch.full((1, 2, 4, 3), float(index + 1)), + values=torch.full((1, 2, 4, 3), float(index + 11)), + ) + for index in range(layer_count) + ], + captured={}, + ) + + @staticmethod + def _inputs(cache_position: int) -> dict[str, torch.Tensor]: + return { + "decoder_input_ids": torch.tensor([[7]]), + "decoder_attention_mask": torch.ones((1, 5), dtype=torch.long), + "encoder_hidden_states": torch.zeros((1, 3, 4)), + "cache_position": torch.tensor(cache_position), + } + + def test_invokes_legacy_decoder_with_empty_cache_at_position_zero(self) -> None: + cache = self._cache() + returned = tuple( + ( + torch.full((1, 2, 1, 3), float(index + 21)), + torch.full((1, 2, 1, 3), float(index + 31)), + ) + for index in range(len(cache.layers)) + ) + received: dict[str, object] = {} + + def decoder(**kwargs: object) -> SimpleNamespace: + received.update(kwargs) + return SimpleNamespace( + last_hidden_state=torch.ones((1, 1, 4)), + past_key_values=returned, + ) + + wrapper = HF_MODEL_CLASS_MAPPING[("florence2", "text2text-generation")]() + wrapper.model = SimpleNamespace( + get_decoder=lambda: decoder, + language_model=SimpleNamespace(lm_head=lambda hidden_states: hidden_states), + ) + + logits = wrapper._invoke_hf(cache, self._inputs(cache_position=0)) + + legacy_cache = received["past_key_values"] + assert isinstance(legacy_cache, tuple) + assert [len(layer) for layer in legacy_cache] == [2] * len(cache.layers) + assert [layer[0].shape[2] for layer in legacy_cache] == [0] * len(cache.layers) + assert received["attention_mask"].shape == (1, 1) + for index, (key, value) in enumerate(returned): + assert torch.equal(cache.captured[index][0], key) + assert torch.equal(cache.captured[index][1], value) + assert torch.equal(logits, torch.ones((1, 1, 4))) + + def test_uses_active_prefix_for_later_cache_position(self) -> None: + cache = self._cache() + returned = tuple( + ( + torch.full((1, 2, 3, 3), float(index + 41)), + torch.full((1, 2, 3, 3), float(index + 51)), + ) + for index in range(len(cache.layers)) + ) + received: dict[str, object] = {} + + def decoder(**kwargs: object) -> SimpleNamespace: + received.update(kwargs) + return SimpleNamespace( + last_hidden_state=torch.ones((1, 1, 4)), + past_key_values=returned, + ) + + wrapper = HF_MODEL_CLASS_MAPPING[("florence2", "text2text-generation")]() + wrapper.model = SimpleNamespace( + get_decoder=lambda: decoder, + language_model=SimpleNamespace(lm_head=lambda hidden_states: hidden_states), + ) + + wrapper._invoke_hf(cache, self._inputs(cache_position=2)) + + legacy_cache = received["past_key_values"] + assert isinstance(legacy_cache, tuple) + assert [layer[0].shape[2] for layer in legacy_cache] == [2] * len(cache.layers) + assert torch.equal(legacy_cache[1][1], cache.layers[1].values[:, :, :2, :]) + assert received["attention_mask"].shape == (1, 3) + for index, (key, value) in enumerate(returned): + assert torch.equal(cache.captured[index][0], key[:, :, 2:, :]) + assert torch.equal(cache.captured[index][1], value[:, :, 2:, :]) + + +class TestFlorence2Recipes: + def test_decoder_recipe_declares_only_dynamic_encoder_axes(self) -> None: + export = _load_recipe("image-to-text_fp16_config_decoder.json")["export"] + + assert export["dynamic_axes"] == { + "encoder_hidden_states": { + "0": "batch_size", + "1": "source_sequence_length", + } + } + + def test_decoder_recipe_assembles_dynamic_source_shape(self) -> None: + config = WinMLBuildConfig.from_dict( + _load_recipe("image-to-text_fp16_config_decoder.json") + ) + + assert config.export is not None + assert config.export.dynamic_axes == { + "encoder_hidden_states": { + 0: "batch_size", + 1: "source_sequence_length", + } + } + assert next( + tensor + for tensor in config.export.input_tensors or [] + if tensor.name == "encoder_hidden_states" + ).shape == (1, 16, 768) + + def test_component_recipes_share_remote_code_consent(self) -> None: + encoder_loader = _load_recipe("image-to-text_fp16_config_encoder.json")["loader"] + decoder_loader = _load_recipe("image-to-text_fp16_config_decoder.json")["loader"] + + assert encoder_loader["trust_remote_code"] is True + assert decoder_loader["trust_remote_code"] is True + + def test_native_loader_and_processor_forward_explicit_consent(self, monkeypatch) -> None: + from winml.modelkit.models.hf.florence2 import ( + _load_native_combined_processor, + _NativeFlorence2ForConditionalGeneration, + ) + + model_kwargs: dict[str, object] = {} + processor_kwargs: dict[str, object] = {} + + def load_model(*args: object, **kwargs: object) -> tuple[object, dict[str, list[str]]]: + model_kwargs.update(kwargs) + return object(), { + "missing_keys": [], + "unexpected_keys": [], + "mismatched_keys": [], + } + + def load_processor(*args: object, **kwargs: object) -> object: + processor_kwargs.update(kwargs) + return object() + + monkeypatch.setattr("transformers.AutoModelForCausalLM.from_pretrained", load_model) + monkeypatch.setattr("transformers.AutoProcessor.from_pretrained", load_processor) + + _NativeFlorence2ForConditionalGeneration.from_pretrained( + "local-model", + trust_remote_code=False, + ) + _load_native_combined_processor("local-model", trust_remote_code=False) + + assert model_kwargs["trust_remote_code"] is False + assert processor_kwargs["trust_remote_code"] is False + + def test_decoder_recipe_matches_base_text_config(self) -> None: + config = _florence2_base_config() + export = _load_recipe("image-to-text_fp16_config_decoder.json")["export"] + tensors = export["input_tensors"] + shapes = {tensor["name"]: tensor["shape"] for tensor in tensors} + value_ranges = { + tensor["name"]: tensor["value_range"] for tensor in tensors if "value_range" in tensor + } + past_key_names = [ + name for name in shapes if name.startswith("past_") and name.endswith("_key") + ] + output_names = [tensor["name"] for tensor in export["output_tensors"]] + + assert shapes["encoder_hidden_states"][2] == config.text_config.d_model + assert shapes["past_0_key"][1] == config.text_config.decoder_attention_heads + assert shapes["past_0_key"][3] == ( + config.text_config.d_model // config.text_config.decoder_attention_heads + ) + assert value_ranges["decoder_input_ids"] == [0, config.text_config.vocab_size] + assert past_key_names == [ + f"past_{index}_key" for index in range(config.text_config.decoder_layers) + ] + assert output_names == [ + "logits", + *( + name + for index in range(config.text_config.decoder_layers) + for name in (f"present_{index}_key", f"present_{index}_value") + ), + ] + + +class TestFlorence2Generation: + def test_prompt_ids_stay_on_encoder_and_decoder_steps_are_single_token(self) -> None: + """Prompt IDs must remain encoder inputs when the primary input is pixels.""" + from winml.modelkit.models.hf.florence2 import WinMLFlorence2ImageToText + + class Encoder: + io_config: ClassVar[dict[str, list[object]]] = { + "input_names": ["input_ids", "pixel_values", "attention_mask"], + "input_shapes": [[1, 8], [1, 3, 2, 2], [1, 8]], + } + + def __call__(self, **kwargs: torch.Tensor) -> BaseModelOutput: + return BaseModelOutput(last_hidden_state=torch.zeros((1, 2, 4))) + + class Decoder: + io_config: ClassVar[dict[str, list[object]]] = { + "input_names": [ + "decoder_input_ids", + "encoder_hidden_states", + "decoder_attention_mask", + "cache_position", + "past_0_key", + "past_0_value", + ], + "input_shapes": [ + [1, 1], + [1, 2, 4], + [1, 4], + [1], + [1, 1, 4, 4], + [1, 1, 4, 4], + ], + "input_types": [int, float, int, int, "float32", "float32"], + } + + def __init__(self) -> None: + self.decoder_input_ids: list[torch.Tensor] = [] + + def __call__(self, **kwargs: torch.Tensor) -> dict[str, torch.Tensor]: + self.decoder_input_ids.append(kwargs["decoder_input_ids"].clone()) + logits = torch.zeros((1, 1, 8)) + logits[:, :, 3] = 1 + return { + "logits": logits, + "present_0_key": torch.zeros((1, 1, 1, 4)), + "present_0_value": torch.zeros((1, 1, 1, 4)), + } + + config = _florence2_base_config() + config.text_config.num_hidden_layers = 1 + config.text_config.decoder_layers = 1 + decoder = Decoder() + model = WinMLFlorence2ImageToText( + {"encoder": Encoder(), "decoder": decoder}, + config, + ) + + model.generate( + pixel_values=torch.zeros((1, 3, 2, 2)), + input_ids=torch.arange(8).reshape(1, 8), + attention_mask=torch.ones((1, 8), dtype=torch.long), + max_new_tokens=2, + eos_token_id=None, + ) + + assert [tuple(ids.shape) for ids in decoder.decoder_input_ids] == [(1, 1), (1, 1)] diff --git a/tests/unit/inference/test_combined_processor_image_to_text.py b/tests/unit/inference/test_combined_processor_image_to_text.py new file mode 100644 index 000000000..66408286e --- /dev/null +++ b/tests/unit/inference/test_combined_processor_image_to_text.py @@ -0,0 +1,137 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Tests for capability-selected combined image/text preprocessing.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any, ClassVar + +import torch +from PIL import Image +from transformers.feature_extraction_utils import BatchFeature + +from winml.modelkit.inference.pipeline import ( + CombinedProcessorImageToTextPipeline, + PipelineCapability, + create_pipeline, +) + + +class RecordingCombinedProcessor: + """Small processor double with the same batch surface as HF processors.""" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + self.tokenizer = RecordingTokenizer() + + def __call__( + self, *, images: Image.Image, text: str, return_tensors: str + ) -> BatchFeature: + self.calls.append( + {"images": images, "text": text, "return_tensors": return_tensors} + ) + return BatchFeature( + { + "input_ids": torch.tensor([[7, 8]], dtype=torch.long), + "pixel_values": torch.ones((1, 3, 2, 2), dtype=torch.float32), + "attention_mask": torch.ones((1, 2), dtype=torch.long), + } + ) + + +class RecordingTokenizer: + """Tokenizer double used by the combined processor.""" + + model_max_length = 1 + + def decode(self, output_ids: Any, *, skip_special_tokens: bool) -> str: + assert skip_special_tokens is True + return f"decoded-{output_ids.tolist()}" + + +class FakePipeline(SimpleNamespace): + """Pipeline double that supports tokenizer adaptation.""" + + def preprocess(self, inputs: Any, **kwargs: Any) -> Any: + return inputs + + +def test_combined_processor_receives_image_and_prompt() -> None: + processor = RecordingCombinedProcessor() + pipe = object.__new__(CombinedProcessorImageToTextPipeline) + pipe.processor = processor + pipe.framework = "pt" + pipe.model = SimpleNamespace(dtype=torch.float32) + image = Image.new("RGB", (2, 2)) + + batch = pipe.preprocess(image, prompt="") + + assert processor.calls == [ + {"images": image, "text": "", "return_tensors": "pt"} + ] + assert set(batch) == {"input_ids", "pixel_values", "attention_mask"} + assert batch["input_ids"].dtype == torch.long + assert batch["pixel_values"].dtype == torch.float32 + + +class CapabilityModel: + """Minimal model surface required by the shared pipeline factory.""" + + pipeline_capabilities = frozenset( + {PipelineCapability.COMBINED_IMAGE_TEXT_PROCESSOR} + ) + io_config: ClassVar[dict[str, list[list[int]]]] = { + "input_shapes": [[1, 2], [1, 3, 2, 2]] + } + processor = RecordingCombinedProcessor() + + @classmethod + def create_combined_processor(cls, model_id: str) -> object: + assert model_id == "local-model" + return cls.processor + + +def test_factory_uses_capability_selected_pipeline(monkeypatch) -> None: + captured: dict[str, Any] = {} + + def fake_pipeline(*args: Any, **kwargs: Any) -> Any: + assert kwargs["tokenizer"] is CapabilityModel.processor.tokenizer + assert not isinstance(kwargs["tokenizer"], str) + assert "feature_extractor" not in kwargs + assert "image_processor" not in kwargs + captured["args"] = args + captured["kwargs"] = kwargs + return FakePipeline(tokenizer=kwargs["tokenizer"], _preprocess_params={}) + + monkeypatch.setattr("transformers.pipeline", fake_pipeline) + + pipe = create_pipeline("image-to-text", CapabilityModel(), "local-model") + + assert captured["kwargs"]["pipeline_class"] is CombinedProcessorImageToTextPipeline + assert captured["kwargs"]["processor"] is CapabilityModel.processor + assert CombinedProcessorImageToTextPipeline.postprocess( + pipe, + [torch.tensor([1, 2])], + ) == [{"generated_text": "decoded-[1, 2]"}] + + +def test_factory_keeps_default_pipeline_without_capability(monkeypatch) -> None: + captured: dict[str, Any] = {} + + def fake_pipeline(*args: Any, **kwargs: Any) -> Any: + captured["kwargs"] = kwargs + return SimpleNamespace(tokenizer=None) + + monkeypatch.setattr("transformers.pipeline", fake_pipeline) + + create_pipeline( + "image-to-text", + SimpleNamespace(io_config={"input_shapes": [[1, 3, 2, 2]]}), + "local-model", + ) + + assert "pipeline_class" not in captured["kwargs"] diff --git a/tests/unit/inspect/test_composite_support.py b/tests/unit/inspect/test_composite_support.py new file mode 100644 index 000000000..6201c5aa0 --- /dev/null +++ b/tests/unit/inspect/test_composite_support.py @@ -0,0 +1,180 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Registry-driven composite inspect support tests.""" + +from __future__ import annotations + +from typing import ClassVar +from unittest.mock import MagicMock + +import pytest + +from winml.modelkit.inspect.types import ( + ExporterInfo, + IOConfigInfo, + LoaderInfo, + ProcessorInfo, + SupportLevel, + WinMLInfo, +) + + +def _supported_loader() -> LoaderInfo: + return LoaderInfo("ComponentModel", "registry", SupportLevel.SUPPORTED) + + +def _supported_exporter() -> ExporterInfo: + return ExporterInfo( + "ComponentOnnxConfig", + "registry", + SupportLevel.SUPPORTED, + ) + + +def test_composite_exporter_aggregates_every_registered_component( + monkeypatch, +) -> None: + """A composite exporter is supported only after every registered component resolves.""" + from winml.modelkit.inspect import resolver + from winml.modelkit.models.winml.composite_model import COMPOSITE_MODEL_REGISTRY + + class Composite: + _SUB_MODEL_CONFIG: ClassVar = {"first": "component-one", "second": "component-two"} + + calls: list[tuple[str, str]] = [] + + def resolve_loader(model_type: str, task: str) -> LoaderInfo: + calls.append(("loader", task)) + return _supported_loader() + + def resolve_exporter( + model_type: str, + task: str, + hf_config: object | None = None, + *, + model_id: str | None = None, + ) -> ExporterInfo: + calls.append(("exporter", task)) + return _supported_exporter() + + monkeypatch.setitem(COMPOSITE_MODEL_REGISTRY, ("test-composite", "pipeline"), Composite) + monkeypatch.setattr(resolver, "resolve_loader", resolve_loader) + monkeypatch.setattr(resolver, "resolve_exporter", resolve_exporter) + + info = resolver.resolve_composite_exporter("test-composite", "pipeline") + + assert info is not None + assert info.support_level is SupportLevel.SUPPORTED + assert info.onnx_config_source == "COMPOSITE_MODEL_REGISTRY" + assert calls == [ + ("loader", "component-one"), + ("exporter", "component-one"), + ("loader", "component-two"), + ("exporter", "component-two"), + ] + + +def test_composite_exporter_is_unsupported_when_a_component_cannot_export( + monkeypatch, +) -> None: + """One unsupported component prevents the registered composite from being supported.""" + from winml.modelkit.inspect import resolver + from winml.modelkit.models.winml.composite_model import COMPOSITE_MODEL_REGISTRY + + class Composite: + _SUB_MODEL_CONFIG: ClassVar = {"first": "component-one", "second": "component-two"} + + unsupported = ExporterInfo(None, "none", SupportLevel.UNSUPPORTED) + monkeypatch.setitem(COMPOSITE_MODEL_REGISTRY, ("test-composite", "pipeline"), Composite) + monkeypatch.setattr(resolver, "resolve_loader", lambda *_: _supported_loader()) + monkeypatch.setattr( + resolver, + "resolve_exporter", + lambda _model_type, task, **_: unsupported + if task == "component-two" + else _supported_exporter(), + ) + + info = resolver.resolve_composite_exporter("test-composite", "pipeline") + + assert info is not None + assert info.support_level is SupportLevel.UNSUPPORTED + + +def test_optimization_only_build_config_is_not_an_exporter(monkeypatch) -> None: + """Legacy exporter lookup must continue to fallback when a build config has no export.""" + from winml.modelkit.config import WinMLBuildConfig + from winml.modelkit.inspect import resolver + from winml.modelkit.optim import WinMLOptimizationConfig + + monkeypatch.setitem( + resolver.MODEL_BUILD_CONFIGS, + "test-optim-only", + WinMLBuildConfig(optim=WinMLOptimizationConfig()), + ) + with monkeypatch.context() as context: + tasks_manager = MagicMock() + tasks_manager.get_exporter_config_constructor.return_value = None + context.setattr( + "optimum.exporters.tasks.TasksManager", + tasks_manager, + ) + info = resolver.resolve_exporter("test-optim-only", "component-one") + + assert info.support_level is SupportLevel.UNSUPPORTED + + +@pytest.mark.parametrize("trust_remote_code", [False, True]) +def test_legacy_inspect_forwards_trust_remote_code_to_hierarchy( + monkeypatch, + trust_remote_code: bool, +) -> None: + """The legacy inspect API must forward trust consent to hierarchy loading.""" + import winml.modelkit.inspect as inspect_module + + hf_config = MagicMock() + hf_config.model_type = "test-composite" + hf_config.architectures = [] + config_loader = MagicMock(return_value=hf_config) + hierarchy_loader = MagicMock(return_value=MagicMock(hf_module_count=1)) + monkeypatch.setattr(inspect_module.AutoConfig, "from_pretrained", config_loader) + monkeypatch.setattr( + "winml.modelkit.inspect.hierarchy.extract_hierarchy", + hierarchy_loader, + ) + monkeypatch.setattr(inspect_module, "resolve_loader", lambda *_: _supported_loader()) + monkeypatch.setattr( + inspect_module, + "resolve_exporter", + lambda *_args, **_kwargs: _supported_exporter(), + ) + monkeypatch.setattr( + inspect_module, + "resolve_winml", + lambda *_: WinMLInfo("Composite", "registry", SupportLevel.SUPPORTED), + ) + monkeypatch.setattr(inspect_module, "resolve_cache", lambda *_: MagicMock()) + monkeypatch.setattr( + inspect_module, "resolve_processor", lambda *_args, **_kwargs: ProcessorInfo() + ) + monkeypatch.setattr( + inspect_module, "resolve_io_config", lambda *_args, **_kwargs: IOConfigInfo() + ) + monkeypatch.setattr(inspect_module, "get_build_config", lambda *_: None) + monkeypatch.setattr(inspect_module, "resolve_composite_info", lambda *_: None) + monkeypatch.setattr( + inspect_module, + "resolve_composite_exporter", + lambda *_args, **_kwargs: None, + ) + + inspect_module.inspect_model( + "test/model", + include_hierarchy=True, + trust_remote_code=trust_remote_code, + ) + + config_loader.assert_called_once_with("test/model", trust_remote_code=trust_remote_code) + hierarchy_loader.assert_called_once_with("test/model", trust_remote_code=trust_remote_code) diff --git a/tests/unit/inspect/test_hierarchy.py b/tests/unit/inspect/test_hierarchy.py new file mode 100644 index 000000000..178695354 --- /dev/null +++ b/tests/unit/inspect/test_hierarchy.py @@ -0,0 +1,50 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for HuggingFace hierarchy extraction.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest +from torch import nn + +from winml.modelkit.inspect.hierarchy import extract_hierarchy + + +@pytest.mark.parametrize( + ("kwargs", "trust_remote_code"), + [ + ({}, False), + ({"trust_remote_code": True}, True), + ], +) +def test_extract_hierarchy_forwards_remote_code_consent( + monkeypatch: pytest.MonkeyPatch, + kwargs: dict[str, bool], + trust_remote_code: bool, +) -> None: + """Every hierarchy load must use the caller's explicit trust setting.""" + from winml.modelkit.inspect import hierarchy + + pretrained_loader = MagicMock(side_effect=OSError("not cached")) + config_loader = MagicMock(return_value=object()) + model_loader = MagicMock(return_value=nn.Linear(1, 1)) + monkeypatch.setattr(hierarchy.AutoModel, "from_pretrained", pretrained_loader) + monkeypatch.setattr(hierarchy.AutoConfig, "from_pretrained", config_loader) + monkeypatch.setattr(hierarchy.AutoModel, "from_config", model_loader) + + extract_hierarchy("test/model", **kwargs) + + pretrained_loader.assert_called_once_with( + "test/model", + trust_remote_code=trust_remote_code, + local_files_only=True, + ) + config_loader.assert_called_once_with("test/model", trust_remote_code=trust_remote_code) + model_loader.assert_called_once_with( + config_loader.return_value, + trust_remote_code=trust_remote_code, + ) diff --git a/tests/unit/onnx/test_shape_inference.py b/tests/unit/onnx/test_shape_inference.py new file mode 100644 index 000000000..ff09425fd --- /dev/null +++ b/tests/unit/onnx/test_shape_inference.py @@ -0,0 +1,119 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Tests for safe symbolic ONNX shape inference.""" + +from __future__ import annotations + +import os +import threading +from pathlib import Path +from queue import Queue +from shutil import rmtree + +import onnx + +from winml.modelkit.onnx import infer_shapes, shape +from winml.modelkit.onnx.shape import _infer_symbolic_shapes_worker + + +def _make_model() -> onnx.ModelProto: + input_ = onnx.helper.make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1]) + output = onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1]) + node = onnx.helper.make_node("Identity", ["input"], ["output"]) + graph = onnx.helper.make_graph([node], "test", [input_], [output]) + return onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 17)]) + + +def _make_large_model() -> onnx.ModelProto: + input_ = onnx.helper.make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1]) + output = onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1]) + nodes = [] + previous = "input" + for index in range(1_000): + current = f"value_{index}" + nodes.append(onnx.helper.make_node("Identity", [previous], [current])) + previous = current + nodes.append(onnx.helper.make_node("Identity", [previous], ["output"])) + graph = onnx.helper.make_graph(nodes, "large_test", [input_], [output]) + return onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 17)]) + + +def test_symbolic_inference_keeps_parent_cwd_usable_while_worker_runs() -> None: + """The parent CWD remains usable while symbolic inference owns its scratch CWD.""" + working_dir = (Path("temp") / "shape-inference-test").resolve() + rmtree(working_dir, ignore_errors=True) + working_dir.mkdir(parents=True) + + original_cwd = Path.cwd() + os.chdir(working_dir) + try: + results: list[onnx.ModelProto] = [] + caller = threading.Thread(target=lambda: results.append(infer_shapes(_make_large_model()))) + caller.start() + + assert caller.is_alive() + assert Path.cwd() == working_dir + parent_sidecar = working_dir / "parent-sidecar" + parent_sidecar.write_bytes(b"usable") + assert parent_sidecar.read_bytes() == b"usable" + + caller.join(timeout=60) + assert not caller.is_alive() + assert results + assert not list(working_dir.glob("winmlcli_shape_*")) + finally: + os.chdir(original_cwd) + rmtree(working_dir) + + +def test_symbolic_inference_uses_system_scratch_when_cwd_is_unwritable(monkeypatch) -> None: + """Worker scratch must not be created in the caller's current directory.""" + original_temporary_directory = shape.tempfile.TemporaryDirectory + created_dirs: list[Path] = [] + working_dir = (Path("temp") / "shape-inference-unwritable-cwd-test").resolve() + rmtree(working_dir, ignore_errors=True) + working_dir.mkdir(parents=True) + + def temporary_directory(*args, **kwargs): + if kwargs.get("dir") == working_dir: + raise PermissionError("caller CWD is unwritable") + temporary_dir = original_temporary_directory(*args, **kwargs) + created_dirs.append(Path(temporary_dir.name)) + return temporary_dir + + monkeypatch.setattr(shape.tempfile, "TemporaryDirectory", temporary_directory) + original_cwd = Path.cwd() + os.chdir(working_dir) + try: + result = shape.infer_symbolic_shapes(_make_model()) + + assert result.graph.name == "test" + assert created_dirs + assert all(not scratch.is_relative_to(working_dir) for scratch in created_dirs) + assert not list(working_dir.glob("winmlcli_shape_*")) + finally: + os.chdir(original_cwd) + rmtree(working_dir) + + +def test_symbolic_worker_reports_serialization_failure() -> None: + """A failed worker must return its error to the parent explicitly.""" + working_dir = (Path("temp") / "shape-inference-failure-test").resolve() + rmtree(working_dir, ignore_errors=True) + working_dir.mkdir(parents=True) + + original_cwd = Path.cwd() + os.chdir(working_dir) + try: + results: Queue[tuple[bool, str]] = Queue() + _infer_symbolic_shapes_worker(b"not an onnx model", str(working_dir), results) + success, error = results.get_nowait() + + assert success is False + assert "DecodeError" in error + finally: + os.chdir(original_cwd) + rmtree(working_dir) diff --git a/tests/unit/optim/test_error_paths.py b/tests/unit/optim/test_error_paths.py index 256f0a860..e4d390eb4 100644 --- a/tests/unit/optim/test_error_paths.py +++ b/tests/unit/optim/test_error_paths.py @@ -324,7 +324,7 @@ def test_symbolic_failure_falls_back_to_onnx(self, simple_model: onnx.ModelProto with ( patch( - "onnxruntime.tools.symbolic_shape_infer.SymbolicShapeInference.infer_shapes", + "winml.modelkit.onnx.shape.infer_symbolic_shapes", side_effect=RuntimeError("SymPy error"), ), patch( @@ -341,7 +341,7 @@ def test_both_inference_failures_returns_original(self, simple_model: onnx.Model with ( patch( - "onnxruntime.tools.symbolic_shape_infer.SymbolicShapeInference.infer_shapes", + "winml.modelkit.onnx.shape.infer_symbolic_shapes", side_effect=RuntimeError("SymPy error"), ), patch( @@ -360,7 +360,7 @@ def test_symbolic_success_skips_onnx(self, simple_model: onnx.ModelProto) -> Non with ( patch( - "onnxruntime.tools.symbolic_shape_infer.SymbolicShapeInference.infer_shapes", + "winml.modelkit.onnx.shape.infer_symbolic_shapes", return_value=simple_model, ), patch("winml.modelkit.onnx.shape.onnx.shape_inference.infer_shapes", mock_onnx), diff --git a/tests/unit/utils/test_config_utils.py b/tests/unit/utils/test_config_utils.py index eb369cde6..5578bb4a7 100644 --- a/tests/unit/utils/test_config_utils.py +++ b/tests/unit/utils/test_config_utils.py @@ -344,6 +344,63 @@ def test_unknown_nested_fields_ignored(self) -> None: class TestMergeConfigListReplacement: """Test that lists are replaced, not merged.""" + def test_default_base_reconstructs_input_tensor_specs(self) -> None: + """Serialized input specs must work when the base list is absent.""" + merged = merge_config( + WinMLExportConfig(), + {"input_tensors": [{"name": "recipe_input", "shape": [1, 4]}]}, + ) + + assert merged.input_tensors == [InputTensorSpec(name="recipe_input", shape=(1, 4))] + + def test_default_base_reconstructs_output_tensor_specs(self) -> None: + """Serialized output specs must work when the base list is absent.""" + merged = merge_config( + WinMLExportConfig(), + {"output_tensors": [{"name": "recipe_output"}]}, + ) + + assert merged.output_tensors == [OutputTensorSpec(name="recipe_output")] + + def test_config_object_override_reconstructs_tensor_specs(self) -> None: + """Serialized config overrides must retain typed tensor specifications.""" + base = WinMLExportConfig( + input_tensors=[InputTensorSpec(name="base_input", shape=(1, 2))], + output_tensors=[OutputTensorSpec(name="base_output")], + ) + override = WinMLExportConfig( + input_tensors=[InputTensorSpec(name="override_input", shape=(1, 3))], + output_tensors=[OutputTensorSpec(name="override_output")], + ) + + merged = merge_config(base, override) + + assert merged.input_tensors == [InputTensorSpec(name="override_input", shape=(1, 3))] + assert merged.output_tensors == [OutputTensorSpec(name="override_output")] + + def test_recipe_override_reconstructs_tensor_specs(self) -> None: + """Recipe tensor lists must be reconstructed when merged into a build config.""" + base = WinMLBuildConfig( + export=WinMLExportConfig( + input_tensors=[InputTensorSpec(name="base_input", shape=(1, 2))], + output_tensors=[OutputTensorSpec(name="base_output")], + ) + ) + + merged = merge_config( + base, + { + "export": { + "input_tensors": [{"name": "recipe_input", "shape": [1, 4]}], + "output_tensors": [{"name": "recipe_output"}], + } + }, + ) + + assert merged.export is not None + assert merged.export.input_tensors == [InputTensorSpec(name="recipe_input", shape=(1, 4))] + assert merged.export.output_tensors == [OutputTensorSpec(name="recipe_output")] + def test_list_replacement_with_spec_objects(self) -> None: """Test that input_tensors list is replaced, not merged (using InputTensorSpec objects).""" base = WinMLExportConfig( @@ -369,13 +426,8 @@ def test_list_replacement_with_spec_objects(self) -> None: def test_list_replacement_with_raw_dicts_on_simple_list(self) -> None: """Test list replacement with raw dicts on a config without validation. - Note: merge_config replaces lists directly without converting dict items - to their proper types. This works for configs that don't validate list - contents in __post_init__. - - For WinMLExportConfig specifically, using raw dicts in input_tensors - will fail because __post_init__ validates that items are InputTensorSpec - objects with 'shape' attribute. Use InputTensorSpec objects instead. + Typed dataclass lists are reconstructed from serialized dictionaries. + This config's list contains strings, so it is replaced directly. """ # Test with WinMLQuantizationConfig's optional list fields base = WinMLQuantizationConfig(