From 242d84a830d0393da33be6c4fa5953fbd288d36c Mon Sep 17 00:00:00 2001 From: Hualiang Xie Date: Mon, 20 Jul 2026 14:40:26 +0800 Subject: [PATCH] Add export override options to winml eval command Bring winml eval to parity with build/perf by adding --shape-config, --input-specs, --export-config, and --dynamic-axes. These shape the ONNX export when eval builds from a HuggingFace model ID and are ignored (with a warning) for pre-built ONNX inputs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/commands/eval.md | 5 + src/winml/modelkit/commands/eval.py | 78 ++++++++ src/winml/modelkit/eval/config.py | 39 ++++ src/winml/modelkit/eval/evaluate.py | 20 +- tests/unit/commands/test_eval.py | 282 +++++++++++++++++++++++++++- tests/unit/eval/test_eval.py | 4 +- 6 files changed, 415 insertions(+), 13 deletions(-) diff --git a/docs/commands/eval.md b/docs/commands/eval.md index f31a5df28..e5f3074d5 100644 --- a/docs/commands/eval.md +++ b/docs/commands/eval.md @@ -22,6 +22,10 @@ $ winml eval [options] | `--precision` | | `TEXT` | `auto` | Precision used when building the model from a HuggingFace ID. One of `auto`, `fp32`, `fp16`, `int8`, `int16`, or a mixed `w{x}a{y}` spec (e.g., `w8a16`). `fp16`/`fp32` skip quantization. **Ignored** when `-m` is a pre-built `.onnx` file — the precision is already baked in. | | `--device` | | choice | `auto` | Target device. Choices: `auto`, `npu`, `gpu`, `cpu`. `auto` selects the best available device. Combined with `--precision`, this drives the build when `-m` is a HuggingFace ID. | | `--ep` / `--execution-provider` | | `TEXT` | — | Target ONNX Runtime execution provider when finer control than `--device` is needed. Full names (e.g., `QNNExecutionProvider`, `OpenVINOExecutionProvider`, `VitisAIExecutionProvider`) and aliases (`qnn`, `ov`/`openvino`, `vitis`/`vitisai`) are accepted. | +| `--shape-config` | | `PATH` | — | JSON shape overrides used while auto-generating a Hugging Face export config, for example `{"height": 480, "width": 480}`. Applies only when `-m` is a HuggingFace ID that eval builds; **ignored for pre-built `.onnx` inputs**. | +| `--input-specs` | | `PATH` | — | JSON input tensor specs to merge into the Hugging Face export config. Symbolic string dimensions infer dynamic axes. **Ignored for pre-built `.onnx` inputs**. | +| `--export-config` | | `PATH` | — | JSON ONNX export config overrides (opset version, constant folding, etc.) to merge into the Hugging Face export config. **Ignored for pre-built `.onnx` inputs**. | +| `--dynamic-axes` | | `PATH` | — | JSON dynamic axes mapping for Hugging Face ONNX export, for example `{"input_ids": {"0": "batch", "1": "sequence"}}`. **Ignored for pre-built `.onnx` inputs**. | | `--dataset` | | `TEXT` | task default | HuggingFace dataset path (e.g., `imagenet-1k`, `nyu-mll/glue`). If omitted, a default dataset is selected based on the task. | | `--dataset-name` | | `TEXT` | — | Dataset configuration name for multi-config datasets. | | `--dataset-revision` | | `TEXT` | — | Git revision (branch, tag, or commit) of the dataset to load. Use `refs/convert/parquet` for HF datasets that are only served via the parquet mirror. | @@ -130,6 +134,7 @@ $ winml eval -m encoder=encoder.onnx -m decoder=decoder.onnx --model-id microsof - **Some dataset requires Hub credentials for gated datasets.** Some datasets (e.g., `imagenet-1k`) require a HuggingFace account with accepted terms of use. Log in with `huggingface-cli login` before running eval on gated data. - **`--shuffle` is on by default.** The random 100-sample slice changes between runs unless you pass `--no-shuffle`. Use `--no-shuffle` when comparing two model variants to ensure they see identical samples. - **`--streaming` skips the local cache.** Streaming mode avoids downloading the full split but prevents random shuffling on large datasets. For reproducible evaluation, download the split once and omit `--streaming`. +- **Export overrides only apply when eval builds from a HuggingFace ID.** `--shape-config`, `--input-specs`, `--export-config`, and `--dynamic-axes` shape the ONNX export that `eval` generates when `-m` is a HuggingFace model ID. When `-m` is a pre-built `.onnx` file, there is no export step, so these flags are ignored and the command prints a warning. - **Column names vary across datasets.** If the evaluator raises a missing-column error, run `winml eval --schema --task ` to inspect the expected schema and use `--column` to remap dataset field names to the expected names. ## See also diff --git a/src/winml/modelkit/commands/eval.py b/src/winml/modelkit/commands/eval.py index f97f72294..e5abaacc9 100644 --- a/src/winml/modelkit/commands/eval.py +++ b/src/winml/modelkit/commands/eval.py @@ -88,6 +88,32 @@ optional_message="Applied during model build. Ignored for pre-built ONNX inputs." ) @cli_utils.max_optim_iterations_option(optional_message="Ignored for pre-built ONNX inputs.") +@cli_utils.shape_config_option( + param_name="shape_config_path", + help_text=( + "JSON with shape overrides for auto-generated HuggingFace export configs. " + "Ignored for pre-built ONNX inputs." + ), +) +@cli_utils.input_specs_option( + help_text=( + "JSON file with input specifications for HuggingFace export. " + "Ignored for pre-built ONNX inputs." + ), +) +@cli_utils.export_config_option( + help_text=( + "ONNX export configuration JSON for HuggingFace model builds " + "(opset_version, do_constant_folding, etc.). Ignored for pre-built ONNX inputs." + ), +) +@cli_utils.dynamic_axes_option( + help_text=( + "JSON dynamic axes mapping for HuggingFace ONNX export " + '(e.g., {"input_ids": {"0": "batch", "1": "sequence"}}). ' + "Ignored for pre-built ONNX inputs." + ), +) @click.option( "--samples", type=int, @@ -180,6 +206,10 @@ def eval( optimize: bool, analyze: bool, max_optim_iterations: int | None, + shape_config_path: Path | None, + input_specs: Path | None, + export_config: Path | None, + dynamic_axes: Path | None, ep: EPNameOrAlias | None, samples: int, split: str, @@ -241,6 +271,7 @@ def eval( # ── 2. Resolve in place ── _resolve_model(cfg, model, model_id) + _apply_export_overrides(cfg, shape_config_path, input_specs, export_config, dynamic_axes) _resolve_device(cfg) _resolve_label_mapping(cfg) _run_dataset_script(cfg, trust_remote_code) @@ -368,6 +399,53 @@ def _resolve_model( cfg.model_id = resolved_id +def _apply_export_overrides( + cfg: WinMLEvaluationConfig, + shape_config_path: Path | None, + input_specs: Path | None, + export_config: Path | None, + dynamic_axes: Path | None, +) -> None: + """Parse the HuggingFace export CLI overrides onto *cfg* (in place). + + ``--shape-config``/``--input-specs``/``--export-config``/``--dynamic-axes`` + only affect the HF-build path (``model_path is None``), where eval exports + and builds the model. A pre-built ONNX input is consumed as-is (no export + step), so any export/shape overrides are dropped with a warning — mirroring + the ``--precision``/build-flag warnings and winml perf's ONNX path. + Requires ``cfg.model_path`` to already be resolved (call after + :func:`_resolve_model`). + """ + export_flags = ( + ("--shape-config", shape_config_path), + ("--input-specs", input_specs), + ("--export-config", export_config), + ("--dynamic-axes", dynamic_axes), + ) + provided = [flag for flag, value in export_flags if value is not None] + if not provided: + return + + if cfg.model_path is not None: + logger.warning( + "%s ignored for pre-built ONNX inputs " + "(no export runs; these apply only when building from a model ID).", + ", ".join(provided), + ) + return + + if shape_config_path is not None: + cfg.shape_config = cli_utils.load_json_object(shape_config_path, "--shape-config") + + export_overrides = cli_utils.load_export_overrides( + export_config=export_config, + input_specs=input_specs, + dynamic_axes=dynamic_axes, + ) + if export_overrides: + cfg.export_overrides = export_overrides + + def _resolve_device(cfg: WinMLEvaluationConfig) -> None: """Resolve ``'auto'`` → concrete device string on *cfg* in place.""" if cfg.device and cfg.device.lower() != "auto": diff --git a/src/winml/modelkit/eval/config.py b/src/winml/modelkit/eval/config.py index 923c0474d..555dbf310 100644 --- a/src/winml/modelkit/eval/config.py +++ b/src/winml/modelkit/eval/config.py @@ -79,6 +79,25 @@ def to_dict(self) -> dict[str, Any]: return result +def _serialize_export_overrides(overrides: dict[str, Any]) -> dict[str, Any]: + """Render export overrides JSON-safe for :meth:`WinMLEvaluationConfig.to_dict`. + + ``--input-specs`` is parsed into ``InputTensorSpec`` objects (via + ``load_export_overrides``), which are not JSON serializable; convert them to + plain dicts. Every other override key (``dynamic_axes`` mapping, opset ints, + bool flags) is already JSON-safe and passed through unchanged. + """ + serialized: dict[str, Any] = {} + for key, value in overrides.items(): + if key == "input_tensors" and isinstance(value, list): + serialized[key] = [ + spec.to_dict() if hasattr(spec, "to_dict") else spec for spec in value + ] + else: + serialized[key] = value + return serialized + + @dataclass class WinMLEvaluationConfig: """Configuration for model evaluation. @@ -92,6 +111,13 @@ class WinMLEvaluationConfig: device: Target device for inference. ep: Explicit execution provider (e.g., "qnn", "dml"). Overrides device-to-provider mapping when provided. + shape_config: Shape overrides for the auto-generated HuggingFace export + config. Only used when building from ``model_id`` (ignored for + pre-built ONNX inputs). + export_overrides: Sparse ONNX export overrides + (``--input-specs``/``--export-config``/``--dynamic-axes``) merged + under the build config's ``export`` section. Only used when building + from ``model_id`` (ignored for pre-built ONNX inputs). dataset: Dataset configuration. output_path: Path to write JSON results. mode: Evaluation mode (see :data:`EvalMode`). @@ -122,6 +148,13 @@ class WinMLEvaluationConfig: optimize: bool = True analyze: bool = True max_optim_iterations: int | None = None + # HuggingFace export overrides, applied only when building from model_id + # (ignored for pre-built ONNX inputs). Shared semantics with winml + # build/perf: ``shape_config`` is passed straight to from_pretrained, while + # ``export_overrides`` (--input-specs/--export-config/--dynamic-axes) is + # merged under the build config's ``export`` section. + shape_config: dict | None = None + export_overrides: dict[str, Any] | None = None dataset: DatasetConfig = field(default_factory=DatasetConfig) output_path: Path | None = field(default=None, metadata={"cli_name": "output"}) mode: EvalMode = "onnx" @@ -153,6 +186,10 @@ def to_dict(self) -> dict: result["analyze"] = self.analyze if self.max_optim_iterations is not None: result["max_optim_iterations"] = self.max_optim_iterations + if self.shape_config: + result["shape_config"] = self.shape_config + if self.export_overrides: + result["export_overrides"] = _serialize_export_overrides(self.export_overrides) result["dataset"] = self.dataset.to_dict() if self.output_path is not None: result["output_path"] = str(self.output_path) @@ -190,6 +227,8 @@ def from_dict(cls, data: dict) -> WinMLEvaluationConfig: optimize=data.get("optimize", True), analyze=data.get("analyze", True), max_optim_iterations=data.get("max_optim_iterations"), + shape_config=data.get("shape_config"), + export_overrides=data.get("export_overrides"), dataset=dataset, output_path=(Path(data["output_path"]) if data.get("output_path") else None), mode=data.get("mode", "onnx"), diff --git a/src/winml/modelkit/eval/evaluate.py b/src/winml/modelkit/eval/evaluate.py index 74863970f..3e9aa3b96 100644 --- a/src/winml/modelkit/eval/evaluate.py +++ b/src/winml/modelkit/eval/evaluate.py @@ -253,7 +253,9 @@ def _load_model(config: WinMLEvaluationConfig) -> WinMLPreTrainedModel | WinMLCo if config.model_path is not None: # Pre-built ONNX: precision is already baked into the model and is - # ignored here (mirrors winml perf's ONNX path). + # ignored here (mirrors winml perf's ONNX path). Export overrides / + # shape_config are HuggingFace-export concepts and never reach here — + # the CLI warns and drops them for pre-built ONNX inputs. from transformers import AutoConfig hf_config = AutoConfig.from_pretrained(config.model_id) @@ -270,6 +272,19 @@ def _load_model(config: WinMLEvaluationConfig) -> WinMLPreTrainedModel | WinMLCo model.config = hf_config return model + # HuggingFace build path — export overrides (--input-specs/--export-config/ + # --dynamic-axes) are merged under the build config's ``export`` section as a + # sparse dict so from_pretrained routes them through merge_export_overrides + # (patching auto-resolved input_tensors by name / re-deriving dynamic_axes). + # Passing a dict rather than a WinMLBuildConfig avoids clobbering the + # auto-resolved export config with default fields. Mirrors winml build/perf. + build_override: Any = quant_override + if config.export_overrides: + override_dict: dict[str, Any] = {"export": config.export_overrides} + if not config.quant: + override_dict["quant"] = None + build_override = override_dict + return WinMLAutoModel.from_pretrained( config.model_id, task=config.task, @@ -277,7 +292,8 @@ def _load_model(config: WinMLEvaluationConfig) -> WinMLPreTrainedModel | WinMLCo precision=config.precision, ep=config.ep, allow_unsupported_nodes=config.allow_unsupported_nodes, - config=quant_override, + config=build_override, + shape_config=config.shape_config, **pipeline_kwargs, ) diff --git a/tests/unit/commands/test_eval.py b/tests/unit/commands/test_eval.py index a73b9f520..43f154064 100644 --- a/tests/unit/commands/test_eval.py +++ b/tests/unit/commands/test_eval.py @@ -234,13 +234,8 @@ def test_composite_hub_refs_resolved(self, tmp_path): enc_local.write_bytes(b"") dec_local = tmp_path / "prompt_encoder_mask_decoder_int8.onnx" dec_local.write_bytes(b"") - enc_ref = ( - "onnx-community/sam3-tracker-ONNX/onnx/vision_encoder_int8.onnx" - ) - dec_ref = ( - "onnx-community/sam3-tracker-ONNX/onnx/" - "prompt_encoder_mask_decoder_int8.onnx" - ) + enc_ref = "onnx-community/sam3-tracker-ONNX/onnx/vision_encoder_int8.onnx" + dec_ref = "onnx-community/sam3-tracker-ONNX/onnx/prompt_encoder_mask_decoder_int8.onnx" # Map each Hub ref to its (different) local cache location. def fake_resolve(ref, **kwargs): @@ -271,9 +266,7 @@ def test_composite_mixed_hub_and_local(self, onnx_vision, tmp_path): """One role is a Hub ref, the other is a local path -- both work.""" dec_local = tmp_path / "decoder.onnx" dec_local.write_bytes(b"") - enc_ref = ( - "onnx-community/sam3-tracker-ONNX/onnx/vision_encoder_int8.onnx" - ) + enc_ref = "onnx-community/sam3-tracker-ONNX/onnx/vision_encoder_int8.onnx" # ``resolve_hf_onnx_path`` is the underlying downloader; the # unified classifier+resolver only calls it for hub_onnx inputs. @@ -1015,3 +1008,272 @@ def test_invalid_format_rejected(self, runner: CliRunner): result = runner.invoke(eval_cmd, ["-m", "test", "--format", "xml"]) assert result.exit_code != 0 + + +# --------------------------------------------------------------------------- +# HuggingFace export overrides (--shape-config/--input-specs/--export-config/ +# --dynamic-axes) — parity with winml build/perf. +# --------------------------------------------------------------------------- + + +class TestEvalExportOverrides: + """Export/shape overrides only affect the HF-build path; ignored for ONNX.""" + + def _write(self, path, payload): + path.write_text(json.dumps(payload), encoding="utf-8") + return path + + def test_help_shows_export_options(self, runner: CliRunner): + from winml.modelkit.commands.eval import eval as eval_cmd + + result = runner.invoke(eval_cmd, ["--help"]) + assert result.exit_code == 0 + for flag in ("--shape-config", "--input-specs", "--export-config", "--dynamic-axes"): + assert flag in result.output + + def test_apply_export_overrides_hf_sets_fields(self, tmp_path): + """HF input: overrides are parsed onto the config.""" + from winml.modelkit.commands.eval import _apply_export_overrides + from winml.modelkit.eval.config import WinMLEvaluationConfig + + input_specs = self._write( + tmp_path / "inputs.json", + {"pixel_values": {"dtype": "float32", "shape": ["batch", 3, 224, 224]}}, + ) + export_config = self._write(tmp_path / "export.json", {"opset_version": 18}) + dynamic_axes = self._write(tmp_path / "da.json", {"pixel_values": {"0": "batch"}}) + shape_config = self._write(tmp_path / "shape.json", {"height": 224, "width": 224}) + + cfg = WinMLEvaluationConfig(model_id="microsoft/resnet-50", model_path=None) + _apply_export_overrides(cfg, shape_config, input_specs, export_config, dynamic_axes) + + assert cfg.shape_config == {"height": 224, "width": 224} + assert cfg.export_overrides is not None + assert cfg.export_overrides["opset_version"] == 18 + assert cfg.export_overrides["dynamic_axes"] == {"pixel_values": {"0": "batch"}} + assert cfg.export_overrides["input_tensors"][0].name == "pixel_values" + assert cfg.export_overrides["input_tensors"][0].shape == ("batch", 3, 224, 224) + + def test_apply_export_overrides_onnx_warns_and_skips(self, tmp_path, caplog): + """Pre-built ONNX input: overrides are dropped with a warning.""" + from winml.modelkit.commands.eval import _apply_export_overrides + from winml.modelkit.eval.config import WinMLEvaluationConfig + + dynamic_axes = self._write(tmp_path / "da.json", {"input_ids": {"0": "batch"}}) + shape_config = self._write(tmp_path / "shape.json", {"height": 224}) + + cfg = WinMLEvaluationConfig(model_id="microsoft/resnet-50", model_path="model.onnx") + with caplog.at_level("WARNING"): + _apply_export_overrides(cfg, shape_config, None, None, dynamic_axes) + + assert cfg.export_overrides is None + assert cfg.shape_config is None + assert "ignored for pre-built ONNX" in caplog.text + assert "--shape-config" in caplog.text + assert "--dynamic-axes" in caplog.text + + def test_apply_export_overrides_none_is_noop(self): + """No flags provided: config stays untouched, no parsing/warnings.""" + from winml.modelkit.commands.eval import _apply_export_overrides + from winml.modelkit.eval.config import WinMLEvaluationConfig + + cfg = WinMLEvaluationConfig(model_id="m", model_path="model.onnx") + _apply_export_overrides(cfg, None, None, None, None) + assert cfg.shape_config is None + assert cfg.export_overrides is None + + def test_load_model_hf_threads_export_overrides(self): + """_load_model forwards export overrides as a sparse {"export": ...} dict.""" + from unittest.mock import MagicMock + + from winml.modelkit.eval.config import WinMLEvaluationConfig + from winml.modelkit.eval.evaluate import _load_model + + cfg = WinMLEvaluationConfig( + model_id="microsoft/resnet-50", + task="image-classification", + device="cpu", + shape_config={"height": 480, "width": 480}, + export_overrides={"dynamic_axes": {"pixel_values": {"0": "batch"}}}, + ) + with patch( + "winml.modelkit.models.auto.WinMLAutoModel.from_pretrained", + return_value=MagicMock(), + ) as mock_fp: + _load_model(cfg) + + kwargs = mock_fp.call_args.kwargs + assert kwargs["config"] == {"export": {"dynamic_axes": {"pixel_values": {"0": "batch"}}}} + assert kwargs["shape_config"] == {"height": 480, "width": 480} + + def test_load_model_hf_export_overrides_with_no_quant(self): + """--no-quant + export overrides fold quant:None into the sparse override.""" + from unittest.mock import MagicMock + + from winml.modelkit.eval.config import WinMLEvaluationConfig + from winml.modelkit.eval.evaluate import _load_model + + cfg = WinMLEvaluationConfig( + model_id="microsoft/resnet-50", + task="image-classification", + device="cpu", + quant=False, + export_overrides={"opset_version": 18}, + ) + with patch( + "winml.modelkit.models.auto.WinMLAutoModel.from_pretrained", + return_value=MagicMock(), + ) as mock_fp: + _load_model(cfg) + + assert mock_fp.call_args.kwargs["config"] == { + "export": {"opset_version": 18}, + "quant": None, + } + + def test_load_model_hf_no_overrides_passes_none(self): + """No overrides (quant default): from_pretrained gets config=None, shape_config=None.""" + from unittest.mock import MagicMock + + from winml.modelkit.eval.config import WinMLEvaluationConfig + from winml.modelkit.eval.evaluate import _load_model + + cfg = WinMLEvaluationConfig( + model_id="microsoft/resnet-50", + task="image-classification", + device="cpu", + ) + with patch( + "winml.modelkit.models.auto.WinMLAutoModel.from_pretrained", + return_value=MagicMock(), + ) as mock_fp: + _load_model(cfg) + + assert mock_fp.call_args.kwargs["config"] is None + assert mock_fp.call_args.kwargs["shape_config"] is None + + def test_cli_hf_forwards_export_overrides(self, runner: CliRunner, tmp_path): + """End-to-end: CLI export flags land on the evaluated config (HF path).""" + from winml.modelkit.commands.eval import eval as eval_cmd + + input_specs = self._write( + tmp_path / "inputs.json", + {"pixel_values": {"dtype": "float32", "shape": ["batch", 3, 224, 224]}}, + ) + export_config = self._write(tmp_path / "export.json", {"opset_version": 18}) + dynamic_axes = self._write(tmp_path / "da.json", {"pixel_values": {"0": "batch"}}) + shape_config = self._write(tmp_path / "shape.json", {"height": 224, "width": 224}) + + captured: dict = {} + + def _fake_evaluate(cfg): + captured["cfg"] = cfg + + class _R: + config = cfg + metrics = {"accuracy": 1.0} # noqa: RUF012 + + def to_dict(self): + return {"metrics": self.metrics, "config": cfg.to_dict()} + + return _R() + + with ( + patch("winml.modelkit.eval.evaluate", side_effect=_fake_evaluate), + patch("winml.modelkit.commands.eval._resolve_device", return_value=None), + patch("winml.modelkit.commands.eval._write_and_display", return_value=None), + ): + result = runner.invoke( + eval_cmd, + [ + "-m", + "microsoft/resnet-50", + "--task", + "image-classification", + "--shape-config", + str(shape_config), + "--input-specs", + str(input_specs), + "--export-config", + str(export_config), + "--dynamic-axes", + str(dynamic_axes), + ], + obj={"debug": False}, + ) + + assert result.exit_code == 0, result.output + cfg = captured["cfg"] + assert cfg.shape_config == {"height": 224, "width": 224} + assert cfg.export_overrides["opset_version"] == 18 + assert cfg.export_overrides["dynamic_axes"] == {"pixel_values": {"0": "batch"}} + assert cfg.export_overrides["input_tensors"][0].name == "pixel_values" + + def test_cli_onnx_ignores_export_overrides(self, runner: CliRunner, tmp_path, onnx_file): + """End-to-end: CLI export flags are dropped for a pre-built ONNX input.""" + from winml.modelkit.commands.eval import eval as eval_cmd + + dynamic_axes = self._write(tmp_path / "da.json", {"pixel_values": {"0": "batch"}}) + + captured: dict = {} + + def _fake_evaluate(cfg): + captured["cfg"] = cfg + + class _R: + config = cfg + metrics = {"accuracy": 1.0} # noqa: RUF012 + + def to_dict(self): + return {"metrics": self.metrics, "config": cfg.to_dict()} + + return _R() + + with ( + patch("winml.modelkit.eval.evaluate", side_effect=_fake_evaluate), + patch("winml.modelkit.commands.eval._resolve_device", return_value=None), + patch("winml.modelkit.commands.eval._write_and_display", return_value=None), + ): + result = runner.invoke( + eval_cmd, + [ + "-m", + str(onnx_file), + "--model-id", + "microsoft/resnet-50", + "--task", + "image-classification", + "--dynamic-axes", + str(dynamic_axes), + ], + obj={"debug": False}, + ) + + assert result.exit_code == 0, result.output + cfg = captured["cfg"] + assert cfg.export_overrides is None + assert cfg.shape_config is None + + def test_to_dict_export_overrides_json_safe(self): + """to_dict serializes InputTensorSpec-bearing overrides to JSON-safe dicts.""" + from winml.modelkit.eval.config import WinMLEvaluationConfig + from winml.modelkit.onnx import InputTensorSpec + + cfg = WinMLEvaluationConfig( + model_id="m", + shape_config={"height": 480}, + export_overrides={ + "opset_version": 18, + "dynamic_axes": {"input_ids": {"0": "batch"}}, + "input_tensors": [ + InputTensorSpec(name="input_ids", dtype="int64", shape=("batch", "seq")) + ], + }, + ) + d = cfg.to_dict() + assert d["shape_config"] == {"height": 480} + # Must round-trip through json.dumps without a TypeError. + dumped = json.loads(json.dumps(d["export_overrides"])) + assert dumped["opset_version"] == 18 + assert dumped["input_tensors"][0]["name"] == "input_ids" + assert dumped["input_tensors"][0]["shape"] == ["batch", "seq"] diff --git a/tests/unit/eval/test_eval.py b/tests/unit/eval/test_eval.py index 1fdbf024e..4a7f35f2a 100644 --- a/tests/unit/eval/test_eval.py +++ b/tests/unit/eval/test_eval.py @@ -1341,7 +1341,8 @@ def test_load_model_from_pretrained(self): # quant defaults to True -> no quant override (config=None); optimize/ # analyze default True with max_optim_iterations=None -> no extra build - # kwargs (build_pipeline_extra_kwargs returns {}). + # kwargs (build_pipeline_extra_kwargs returns {}). shape_config defaults + # to None (no --shape-config given). mock_auto.from_pretrained.assert_called_once_with( "test/model", task="image-classification", @@ -1350,6 +1351,7 @@ def test_load_model_from_pretrained(self): ep=None, allow_unsupported_nodes=False, config=None, + shape_config=None, ) assert result is mock_model