diff --git a/docs/commands/eval.md b/docs/commands/eval.md index 7e5c80ab9..a6c316bbb 100644 --- a/docs/commands/eval.md +++ b/docs/commands/eval.md @@ -1,6 +1,6 @@ # winml eval -> Evaluate ONNX model accuracy on a standard dataset. +> Evaluate ONNX or native Hugging Face PyTorch model accuracy on a standard dataset. ## When to use this @@ -26,6 +26,7 @@ $ winml eval [options] | `--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**. | +| `--runtime` | | `winml\|pytorch` | `winml` | Evaluation runtime. `winml` exports Hugging Face checkpoints to ONNX; `pytorch` evaluates the original checkpoint and supports `auto`, `cpu`, or CUDA-backed `gpu` devices. | | `--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. | @@ -45,7 +46,9 @@ $ winml eval [options] ## How it works -`winml eval` loads the model and runs the evaluation pipeline via the internal `evaluate` function (supporting both HuggingFace IDs and local ONNX files), then pulls the requested number of samples from a HuggingFace dataset. Each sample is preprocessed using the tokenizer or image processor associated with the model ID, passed through the ONNX Runtime session, and the output is compared against the ground-truth label. Aggregated metrics (accuracy, F1, etc.) are printed to the console and optionally written to a JSON file. When `-m` is an ONNX file, `--model-id` must be provided so the command knows which preprocessor and label vocabulary to use. +`winml eval` loads the model and runs the evaluation pipeline via the internal `evaluate` function, then pulls the requested number of samples from a HuggingFace dataset. By default, Hugging Face model IDs and local checkpoints use the `winml` runtime: they are exported to ONNX and evaluated through WinML. With `--runtime pytorch`, the checkpoint-declared PyTorch class and stored dtype are preserved and the same dataset preprocessing, Hugging Face pipeline, task evaluator, and metrics run directly against that model. PyTorch `auto` selects CUDA when available and otherwise CPU; `gpu` requires CUDA. The JSON report identifies the effective runtime as `winml` or `pytorch`. + +Pre-built ONNX files, composite `role=path` models, GenAI bundles, text-generation evaluation, compare mode, references, and tensor input archives remain on their existing WinML paths. `--runtime pytorch` rejects those forms, along with ONNX build, export, EP, precision, quantization, optimization, analysis, and cache-related options. ## Examples @@ -55,6 +58,12 @@ Evaluate a HuggingFace model using the task-default dataset: $ winml eval -m microsoft/resnet-50 ``` +Evaluate the original Hugging Face PyTorch checkpoint on CPU without ONNX export: + +```bash +$ winml eval -m microsoft/resnet-50 --runtime pytorch --device cpu +``` + ```text Task: image-classification Dataset: timm/mini-imagenet (test, 100 samples) @@ -163,6 +172,7 @@ model-build pipeline from the runtime compilation cache. - **`--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. +- **The PyTorch runtime accepts only Hugging Face checkpoints.** `--runtime pytorch` cannot be combined with ONNX files, composite models, GenAI bundles, compare/reference/input-data modes, EP selection, or ONNX build/export controls. Use `--device cpu`, `--device gpu` with CUDA, or `--device auto`. - **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 7c06dd798..8fcd41ebb 100644 --- a/src/winml/modelkit/commands/eval.py +++ b/src/winml/modelkit/commands/eval.py @@ -22,7 +22,7 @@ if TYPE_CHECKING: - from ..eval import EvalResult, WinMLEvaluationConfig + from ..eval import EvalResult, EvalRuntime, WinMLEvaluationConfig from ..utils.constants import EPNameOrAlias @@ -115,6 +115,14 @@ "Ignored for pre-built ONNX inputs." ), ) +@click.option( + "--runtime", + type=click.Choice(["winml", "pytorch"]), + default="winml", + show_default=True, + help="Evaluation runtime. 'winml' exports Hugging Face checkpoints to ONNX; " + "'pytorch' evaluates the original checkpoint.", +) @click.option( "--samples", type=int, @@ -166,7 +174,9 @@ default=None, help="Path to a Python script that builds the evaluation dataset.", ) -@cli_utils.trust_remote_code_option(optional_message="Required when --dataset-script is used.") +@cli_utils.trust_remote_code_option( + optional_message="Used for native Hugging Face loading and required with --dataset-script." +) @cli_utils.allow_unsupported_nodes_option() @click.option( "--schema", @@ -235,6 +245,7 @@ def eval( input_specs: Path | None, export_config: Path | None, dynamic_axes: Path | None, + runtime: EvalRuntime, ep: EPNameOrAlias | None, samples: int, split: str, @@ -302,6 +313,13 @@ def eval( # ── 1. Build config: defaults ← config file ← CLI ── cfg, config_fields = _build_eval_config(ctx, config_file, column, label_mapping_path) + if cfg.runtime not in ("winml", "pytorch"): + raise click.UsageError( + f"Invalid eval runtime {cfg.runtime!r}; expected 'winml' or 'pytorch'." + ) + if cfg.runtime == "pytorch": + _validate_pytorch_runtime_options(ctx, cfg, config_fields) + if cfg.input_data is not None and cfg.mode != "compare": raise click.UsageError("--input-data is only valid with --mode compare.") @@ -310,12 +328,25 @@ def eval( # ── 2. Resolve in place ── _resolve_model(cfg, model, model_id, allow_missing_model_id=cfg.reference_path is not None) + if cfg.runtime == "pytorch" and cfg.model_path is not None: + raise click.UsageError( + "--runtime pytorch requires a Hugging Face model ID or local Hugging Face " + "checkpoint; ONNX files, composite role=path models, and GenAI bundles " + "are not supported." + ) + if cfg.runtime == "pytorch": + from ..eval.evaluate import _validate_pytorch_runtime_config + + try: + _validate_pytorch_runtime_config(cfg) + except ValueError as error: + raise click.UsageError(str(error)) from error _resolve_reference(cfg) _apply_export_overrides(cfg, shape_config_path, input_specs, export_config, dynamic_axes) _resolve_device(cfg) _resolve_genai_ep(ctx, cfg) _resolve_label_mapping(cfg) - _run_dataset_script(cfg, trust_remote_code) + _run_dataset_script(cfg, cfg.trust_remote_code) # Refuse to clobber an existing report unless the user opted in — fail fast # before the (expensive) evaluation runs. @@ -383,7 +414,12 @@ def _build_eval_config( # ── Config file layer (only explicitly-present keys) ── if config_file is not None: - _, raw = cli_utils.load_build_config(config_file) + from ..eval.config import _UnsupportedEvalRuntimeFieldError + + try: + _, raw = cli_utils.load_build_config(config_file) + except _UnsupportedEvalRuntimeFieldError as error: + raise click.UsageError(str(error)) from error # Loader task as lowest-priority fallback loader_section = raw.get("loader") or {} @@ -398,6 +434,14 @@ def _build_eval_config( # Eval section overrides loader/compile fallbacks eval_data = raw.get("eval") if eval_data: + eval_data = dict(eval_data) + legacy_runtime_fields = {"backend", "export_model"}.intersection(eval_data) + if legacy_runtime_fields: + fields = ", ".join(sorted(legacy_runtime_fields)) + raise click.UsageError( + f"Unsupported eval runtime field(s): {fields}. Use 'runtime' " + "with 'winml' or 'pytorch' instead." + ) config_fields.update(eval_data) cfg = merge_config(cfg, eval_data) @@ -442,6 +486,8 @@ def _model_build_bypass(cfg: WinMLEvaluationConfig) -> _ModelBuildBypass | None: from ..eval.evaluate import _ModelLoaderKind, _select_model_loader loader = _select_model_loader(cfg) + if loader is _ModelLoaderKind.PYTORCH: + return _ModelBuildBypass("PyTorch runtime evaluation") if loader is _ModelLoaderKind.GENAI: return _ModelBuildBypass( reason="GenAI bundles", @@ -512,6 +558,71 @@ def _resolve_model_loader_task(cfg: WinMLEvaluationConfig) -> None: cfg.task = _infer_task(cfg) +_PYTORCH_RUNTIME_INCOMPATIBLE_OPTIONS: dict[str, str] = { + "mode": "--mode", + "input_data": "--input-data", + "reference": "--reference", + "ep": "--ep", + "precision": "--precision", + "quant": "--quant/--no-quant", + "optimize": "--optimize/--no-optimize", + "analyze": "--analyze/--no-analyze", + "max_optim_iterations": "--max-optim-iterations", + "shape_config_path": "--shape-config", + "input_specs": "--input-specs", + "export_config": "--export-config", + "dynamic_axes": "--dynamic-axes", + "allow_unsupported_nodes": "--allow-unsupported-nodes", + "skip_build": "--skip-build/--no-skip-build", + "use_cache": "--use-cache/--no-use-cache", + "rebuild": "--rebuild/--no-rebuild", +} + +_PYTORCH_RUNTIME_INCOMPATIBLE_CONFIG_FIELDS = { + "mode", + "input_data", + "reference_path", + "ep", + "precision", + "quant", + "optimize", + "analyze", + "max_optim_iterations", + "shape_config", + "export_overrides", + "allow_unsupported_nodes", + "skip_build", + "use_cache", + "rebuild", +} + + +def _validate_pytorch_runtime_options( + ctx: click.Context, + cfg: WinMLEvaluationConfig, + config_fields: set[str], +) -> None: + """Reject options whose semantics require ONNX export or ONNX Runtime.""" + if cfg.device.lower() not in ("auto", "cpu", "gpu"): + raise click.UsageError( + f"--device {cfg.device} is not supported with --runtime pytorch; use auto, cpu, or gpu." + ) + incompatible = [ + flag + for param_name, flag in _PYTORCH_RUNTIME_INCOMPATIBLE_OPTIONS.items() + if cli_utils.is_cli_provided(ctx, param_name) + ] + incompatible.extend( + f"eval.{field}" + for field in sorted(config_fields & _PYTORCH_RUNTIME_INCOMPATIBLE_CONFIG_FIELDS) + ) + if incompatible: + raise click.UsageError( + "--runtime pytorch cannot be combined with incompatible options: " + f"{', '.join(incompatible)}." + ) + + def _resolve_model( cfg: WinMLEvaluationConfig, model: tuple[str, ...], @@ -520,6 +631,8 @@ def _resolve_model( allow_missing_model_id: bool = False, ) -> None: """Resolve ``-m`` / ``--model-id`` into ``cfg.model_path`` / ``cfg.model_id``.""" + if not model and model_id is None and (cfg.model_path is not None or cfg.model_id is not None): + return model_path, resolved_id = _resolve_model_path( model=model, model_id=model_id, allow_missing_model_id=allow_missing_model_id ) @@ -620,6 +733,17 @@ def _apply_export_overrides( def _resolve_device(cfg: WinMLEvaluationConfig) -> None: """Resolve ``'auto'`` → concrete device string on *cfg* in place.""" + if cfg.runtime == "pytorch": + from ..loader import resolve_native_device + + try: + resolved = resolve_native_device(cfg.device) + except ValueError as error: + raise click.UsageError(str(error)) from error + cfg._auto_device_selected = cfg.device.lower() == "auto" + cfg.device = resolved.name + return + if cfg.device and cfg.device.lower() != "auto": return @@ -888,6 +1012,7 @@ def display_eval_report(result: EvalResult, console: Console) -> None: # Info section console.print() console.print(f"[dim]Task:[/dim] {cfg.task}") + console.print(f"[dim]Runtime:[/dim] {cfg.runtime}") console.print(f"[dim]Device:[/dim] {cfg.device}") if cfg.input_data: console.print(f"[dim]Input data:[/dim] {cfg.input_data}") diff --git a/src/winml/modelkit/eval/__init__.py b/src/winml/modelkit/eval/__init__.py index 5d7f8be5a..435601a15 100644 --- a/src/winml/modelkit/eval/__init__.py +++ b/src/winml/modelkit/eval/__init__.py @@ -14,7 +14,7 @@ from typing import TYPE_CHECKING, Any from .base_evaluator import WinMLEvaluator -from .config import DatasetConfig, WinMLEvaluationConfig +from .config import DatasetConfig, EvalRuntime, WinMLEvaluationConfig from .evaluate import EvalResult, evaluate, get_evaluator_class @@ -47,57 +47,50 @@ _LAZY_ATTRS: dict[str, str] = { # Evaluators - "WinMLDepthEstimationEvaluator": - ".depth_estimation_evaluator:WinMLDepthEstimationEvaluator", - "WinMLFeatureExtractionEvaluator": - ".feature_extraction_evaluator:WinMLFeatureExtractionEvaluator", - "WinMLFillMaskEvaluator": - ".fill_mask_evaluator:WinMLFillMaskEvaluator", - "WinMLImageFeatureExtractionEvaluator": - ".image_feature_extraction_evaluator:WinMLImageFeatureExtractionEvaluator", - "WinMLImageSegmentationEvaluator": - ".image_segmentation_evaluator:WinMLImageSegmentationEvaluator", - "WinMLImageToTextEvaluator": - ".image_to_text_evaluator:WinMLImageToTextEvaluator", - "WinMLKeypointDetectionEvaluator": - ".keypoint_detection_evaluator:WinMLKeypointDetectionEvaluator", - "WinMLObjectDetectionEvaluator": - ".object_detection_evaluator:WinMLObjectDetectionEvaluator", - "WinMLQuestionAnsweringEvaluator": - ".question_answering_evaluator:WinMLQuestionAnsweringEvaluator", - "WinMLTextClassificationEvaluator": - ".text_classification_evaluator:WinMLTextClassificationEvaluator", - "WinMLTextGenerationEvaluator": - ".text_generation_evaluator:WinMLTextGenerationEvaluator", - "WinMLTokenClassificationEvaluator": - ".token_classification_evaluator:WinMLTokenClassificationEvaluator", - "WinMLZeroShotClassificationEvaluator": - ".zero_shot_classification_evaluator:WinMLZeroShotClassificationEvaluator", - "WinMLZeroShotImageClassificationEvaluator": - ".zero_shot_image_classification_evaluator:WinMLZeroShotImageClassificationEvaluator", - "TensorSimilarityEvaluator": - ".tensor_similarity_evaluator:TensorSimilarityEvaluator", + "WinMLDepthEstimationEvaluator": ".depth_estimation_evaluator:WinMLDepthEstimationEvaluator", + "WinMLFeatureExtractionEvaluator": ( + ".feature_extraction_evaluator:WinMLFeatureExtractionEvaluator" + ), + "WinMLFillMaskEvaluator": ".fill_mask_evaluator:WinMLFillMaskEvaluator", + "WinMLImageFeatureExtractionEvaluator": ( + ".image_feature_extraction_evaluator:WinMLImageFeatureExtractionEvaluator" + ), + "WinMLImageSegmentationEvaluator": ( + ".image_segmentation_evaluator:WinMLImageSegmentationEvaluator" + ), + "WinMLImageToTextEvaluator": ".image_to_text_evaluator:WinMLImageToTextEvaluator", + "WinMLKeypointDetectionEvaluator": ( + ".keypoint_detection_evaluator:WinMLKeypointDetectionEvaluator" + ), + "WinMLObjectDetectionEvaluator": ".object_detection_evaluator:WinMLObjectDetectionEvaluator", + "WinMLQuestionAnsweringEvaluator": ( + ".question_answering_evaluator:WinMLQuestionAnsweringEvaluator" + ), + "WinMLTextClassificationEvaluator": ( + ".text_classification_evaluator:WinMLTextClassificationEvaluator" + ), + "WinMLTextGenerationEvaluator": ".text_generation_evaluator:WinMLTextGenerationEvaluator", + "WinMLTokenClassificationEvaluator": ( + ".token_classification_evaluator:WinMLTokenClassificationEvaluator" + ), + "WinMLZeroShotClassificationEvaluator": ( + ".zero_shot_classification_evaluator:WinMLZeroShotClassificationEvaluator" + ), + "WinMLZeroShotImageClassificationEvaluator": ( + ".zero_shot_image_classification_evaluator:WinMLZeroShotImageClassificationEvaluator" + ), + "TensorSimilarityEvaluator": ".tensor_similarity_evaluator:TensorSimilarityEvaluator", # Metrics (defer numpy / scipy / torch / torchmetrics until first use) - "ClassificationMetric": - ".metrics.classification:ClassificationMetric", - "DepthMetric": - ".metrics.depth:DepthMetric", - "KeypointAPMetric": - ".metrics.keypoint:KeypointAPMetric", - "IGNORE_INDEX": - ".metrics.mean_iou:IGNORE_INDEX", - "KNNAccuracyMetric": - ".metrics.knn_accuracy:KNNAccuracyMetric", - "MAPMetric": - ".metrics.mean_average_precision:MAPMetric", - "MeanIoUMetric": - ".metrics.mean_iou:MeanIoUMetric", - "PseudoPerplexityMetric": - ".metrics.pseudo_perplexity:PseudoPerplexityMetric", - "SpearmanCorrelationMetric": - ".metrics.spearman_correlation:SpearmanCorrelationMetric", - "TopKAccuracyMetric": - ".metrics.top_k_accuracy:TopKAccuracyMetric", + "ClassificationMetric": ".metrics.classification:ClassificationMetric", + "DepthMetric": ".metrics.depth:DepthMetric", + "KeypointAPMetric": ".metrics.keypoint:KeypointAPMetric", + "IGNORE_INDEX": ".metrics.mean_iou:IGNORE_INDEX", + "KNNAccuracyMetric": ".metrics.knn_accuracy:KNNAccuracyMetric", + "MAPMetric": ".metrics.mean_average_precision:MAPMetric", + "MeanIoUMetric": ".metrics.mean_iou:MeanIoUMetric", + "PseudoPerplexityMetric": ".metrics.pseudo_perplexity:PseudoPerplexityMetric", + "SpearmanCorrelationMetric": ".metrics.spearman_correlation:SpearmanCorrelationMetric", + "TopKAccuracyMetric": ".metrics.top_k_accuracy:TopKAccuracyMetric", } @@ -124,6 +117,7 @@ def __dir__() -> list[str]: "DatasetConfig", "DepthMetric", "EvalResult", + "EvalRuntime", "KNNAccuracyMetric", "KeypointAPMetric", "MAPMetric", diff --git a/src/winml/modelkit/eval/base_evaluator.py b/src/winml/modelkit/eval/base_evaluator.py index fa2077123..fa5de9dd5 100644 --- a/src/winml/modelkit/eval/base_evaluator.py +++ b/src/winml/modelkit/eval/base_evaluator.py @@ -17,8 +17,6 @@ from datasets import Dataset from transformers.pipelines.base import Pipeline - from ..models.winml.base import WinMLPreTrainedModel - from ..models.winml.composite_model import WinMLCompositeModel from .config import DatasetConfig, WinMLEvaluationConfig logger = logging.getLogger(__name__) @@ -44,10 +42,12 @@ class TFPreTrainedModel: class WinMLEvaluator: """Base evaluator. Loads dataset, creates pipeline, runs HF evaluator.""" + supports_native = True + def __init__( self, config: WinMLEvaluationConfig, - model: WinMLPreTrainedModel | WinMLCompositeModel, + model: Any, ) -> None: self.model = model self.config = config @@ -150,7 +150,13 @@ def prepare_pipeline(self) -> Pipeline: assert self.config.task is not None, "config.task is required to build pipeline" return cast( "Pipeline", - create_pipeline(self.config.task, self.model, self.config.model_id), + create_pipeline( + self.config.task, + self.model, + self.config.model_id, + device=self.config.pipeline_device, + trust_remote_code=self.config.trust_remote_code, + ), ) def _fixed_seq_length(self) -> int | None: diff --git a/src/winml/modelkit/eval/config.py b/src/winml/modelkit/eval/config.py index d8b4e8d66..c6cbc4138 100644 --- a/src/winml/modelkit/eval/config.py +++ b/src/winml/modelkit/eval/config.py @@ -9,12 +9,19 @@ from dataclasses import dataclass, field from pathlib import Path -from typing import Any +from typing import Any, Literal from ..utils.constants import EPNameOrAlias from ..utils.eval_utils import EvalMode +EvalRuntime = Literal["winml", "pytorch"] + + +class _UnsupportedEvalRuntimeFieldError(ValueError): + """Raised when a removed eval runtime field is deserialized.""" + + @dataclass class DatasetConfig: """Dataset configuration, aligned with HF load_dataset() API. @@ -130,6 +137,11 @@ class WinMLEvaluationConfig: from ``model_id`` (ignored for pre-built ONNX inputs). dataset: Dataset configuration. output_path: Path to write JSON results. + runtime: Evaluation runtime. + + - ``"winml"`` (default): export Hugging Face checkpoints to ONNX + and evaluate with WinML. + - ``"pytorch"``: evaluate the original Hugging Face checkpoint. mode: Evaluation mode (see :data:`EvalMode`). - ``"onnx"`` (default): evaluate the ONNX candidate on the @@ -174,11 +186,20 @@ class WinMLEvaluationConfig: skip_build: bool = True use_cache: bool = True rebuild: bool = False + runtime: EvalRuntime = "winml" + trust_remote_code: bool = False _auto_device_selected: bool = field(default=False, repr=False, compare=False, kw_only=True) + @property + def pipeline_device(self) -> str: + """Return the tensor-placement device expected by Transformers pipelines.""" + if self.runtime == "pytorch" and self.device.lower() == "gpu": + return "cuda" + return "cpu" + def to_dict(self) -> dict: """Convert to dictionary for serialization.""" - result: dict = {} + result: dict = {"runtime": self.runtime} if self.model_id is not None: result["model_id"] = self.model_id if self.model_path is not None: @@ -215,14 +236,25 @@ def to_dict(self) -> dict: result["output_path"] = str(self.output_path) if self.mode != "onnx": result["mode"] = self.mode - result["skip_build"] = self.skip_build - result["use_cache"] = self.use_cache - result["rebuild"] = self.rebuild + if self.runtime == "winml": + result["skip_build"] = self.skip_build + result["use_cache"] = self.use_cache + result["rebuild"] = self.rebuild + if self.trust_remote_code: + result["trust_remote_code"] = True return result @classmethod def from_dict(cls, data: dict) -> WinMLEvaluationConfig: """Create from dictionary, ignoring unknown fields.""" + legacy_runtime_fields = {"backend", "export_model"}.intersection(data) + if legacy_runtime_fields: + fields = ", ".join(sorted(legacy_runtime_fields)) + raise _UnsupportedEvalRuntimeFieldError( + f"Unsupported eval runtime field(s): {fields}. Use 'runtime' with " + "'winml' or 'pytorch' instead." + ) + ds_data = data.get("dataset", {}) dataset = DatasetConfig( path=ds_data.get("path"), @@ -259,4 +291,6 @@ def from_dict(cls, data: dict) -> WinMLEvaluationConfig: skip_build=data.get("skip_build", True), use_cache=data.get("use_cache", True), rebuild=data.get("rebuild", False), + runtime=data.get("runtime", "winml"), + trust_remote_code=data.get("trust_remote_code", False), ) diff --git a/src/winml/modelkit/eval/evaluate.py b/src/winml/modelkit/eval/evaluate.py index 3b11238fa..bc85b90e8 100644 --- a/src/winml/modelkit/eval/evaluate.py +++ b/src/winml/modelkit/eval/evaluate.py @@ -22,6 +22,8 @@ if TYPE_CHECKING: from pathlib import Path + from torch import nn + from ..models.winml.base import WinMLPreTrainedModel from ..models.winml.composite_model import WinMLCompositeModel from ..models.winml.genai_causal_lm import WinMLGenaiCausalLM @@ -31,6 +33,7 @@ class _ModelLoaderKind(Enum): + PYTORCH = auto() GENAI = auto() DIRECT_ONNX_COMPARE = auto() EVALUATOR_MANAGED = auto() @@ -40,6 +43,8 @@ class _ModelLoaderKind(Enum): def _select_model_loader(config: WinMLEvaluationConfig) -> _ModelLoaderKind: """Select the model-loading path shared by loading and CLI diagnostics.""" + if config.runtime == "pytorch": + return _ModelLoaderKind.PYTORCH if config.task == "text-generation": return _ModelLoaderKind.GENAI if config.mode == "compare" and config.reference_path is not None: @@ -115,6 +120,59 @@ def get_evaluator_class(config: WinMLEvaluationConfig) -> type[WinMLEvaluator]: return cast("type[WinMLEvaluator]", getattr(module, class_name)) +def _validate_pytorch_runtime_config(config: WinMLEvaluationConfig) -> None: + """Validate state that cannot apply to the PyTorch runtime.""" + if config.runtime == "winml": + return + + incompatible: list[str] = [] + mode = config.mode if config.mode is not None else "onnx" + if mode != "onnx": + incompatible.append("mode") + if config.model_path is not None: + incompatible.append("model_path") + if config.input_data is not None: + incompatible.append("input_data") + if config.reference_path is not None: + incompatible.append("reference_path") + if config.ep is not None: + incompatible.append("ep") + if config.precision != "auto": + incompatible.append("precision") + if not config.quant: + incompatible.append("quant") + if not config.optimize: + incompatible.append("optimize") + if not config.analyze: + incompatible.append("analyze") + if config.max_optim_iterations is not None: + incompatible.append("max_optim_iterations") + if config.shape_config is not None: + incompatible.append("shape_config") + if config.export_overrides is not None: + incompatible.append("export_overrides") + if config.allow_unsupported_nodes: + incompatible.append("allow_unsupported_nodes") + if not config.skip_build: + incompatible.append("skip_build") + if not config.use_cache: + incompatible.append("use_cache") + if config.rebuild: + incompatible.append("rebuild") + if incompatible: + raise ValueError( + f"The PyTorch runtime cannot use WinML-only configuration: {', '.join(incompatible)}." + ) + + if config.task is not None: + evaluator_class = get_evaluator_class(config) + if not evaluator_class.supports_native: + raise ValueError( + f"Task '{config.task}' does not use the standard labeled Hugging Face " + "pipeline and is not supported with --runtime pytorch." + ) + + _FE_DEFAULT = { "path": "mteb/stsbenchmark-sts", "split": "test", @@ -264,7 +322,7 @@ def to_dict(self) -> dict[str, Any]: def _load_model( config: WinMLEvaluationConfig, -) -> WinMLPreTrainedModel | WinMLCompositeModel | WinMLGenaiCausalLM | None: +) -> nn.Module | WinMLPreTrainedModel | WinMLCompositeModel | WinMLGenaiCausalLM | None: """Load model from ONNX path or HF model ID. For evaluators that handle their own ORT session construction from a @@ -280,6 +338,20 @@ def _load_model( from ..utils import cli as cli_utils loader = _select_model_loader(config) + if loader is _ModelLoaderKind.PYTORCH: + if config.model_id is None: + raise ValueError("model_id is required for native Hugging Face evaluation.") + from ..loader import load_native_hf_model + + loaded = load_native_hf_model( + config.model_id, + task=config.task, + device=config.device, + trust_remote_code=config.trust_remote_code, + ) + config.device = loaded.device.name + return loaded.model + if loader is _ModelLoaderKind.GENAI: return _load_genai_causal_lm(config) @@ -329,7 +401,11 @@ def _load_model( from ..loader import load_hf_config - hf_config = load_hf_config(AutoConfig, config.model_id) + hf_config = load_hf_config( + AutoConfig, + config.model_id, + trust_remote_code=config.trust_remote_code, + ) model = WinMLAutoModel.from_onnx( # ``model_path`` is narrowed to ``str | dict[str, str]`` here; # cast bridges dict value-type invariance (str vs str | Path). @@ -468,7 +544,11 @@ def _infer_task(config: WinMLEvaluationConfig) -> str: from ..loader import load_hf_config from ..loader.resolution import resolve_task - hf_config = load_hf_config(AutoConfig, config.model_id) + hf_config = load_hf_config( + AutoConfig, + config.model_id, + trust_remote_code=config.trust_remote_code, + ) return resolve_task(hf_config).task @@ -481,6 +561,8 @@ def evaluate(config: WinMLEvaluationConfig) -> EvalResult: """ from ..utils.eval_utils import EVAL_MODES + if config.runtime not in ("winml", "pytorch"): + raise ValueError(f"Invalid runtime {config.runtime!r}; expected 'winml' or 'pytorch'.") mode = config.mode if config.mode is not None else "onnx" if mode not in EVAL_MODES: raise ValueError(f"Invalid mode {mode!r}; expected one of {EVAL_MODES} or None.") @@ -493,6 +575,7 @@ def evaluate(config: WinMLEvaluationConfig) -> EvalResult: task=config.task if onnx_compare else _resolve_task(config), dataset=deepcopy(config.dataset), ) + _validate_pytorch_runtime_config(config) if config.mode != "compare" and config.dataset.path is None: default = _DEFAULT_DATASETS.get(config.task) if config.task is not None else None if default is None: @@ -536,11 +619,7 @@ def evaluate(config: WinMLEvaluationConfig) -> EvalResult: cls = get_evaluator_class(config) try: console.print("[bold]Loading dataset and evaluating...[/bold]") - # ``model`` is ``None`` for composite evaluators that load ORT - # sessions directly from ``config.model_path`` (currently only - # mask-generation). Type-checker can't follow the per-task - # invariant, so suppress here at the unified call site. - task_evaluator = cls(config, model) # type: ignore[arg-type] + task_evaluator = cls(config, model) metrics = task_evaluator.compute() except DatasetValidationError as error: raise ValueError( @@ -581,10 +660,12 @@ def print_config(config: WinMLEvaluationConfig) -> None: output_console.print(f"[bold blue]Reference:[/bold blue] {config.reference_path}") if config.task is not None: output_console.print(f"[bold blue]Task:[/bold blue] {config.task}") + output_console.print(f"[bold blue]Runtime:[/bold blue] {config.runtime}") output_console.print(f"[bold blue]Device:[/bold blue] {config.device}") if config.ep is not None: output_console.print(f"[bold blue]EP:[/bold blue] {config.ep}") - output_console.print(f"[bold blue]Precision:[/bold blue] {config.precision}") + if config.runtime == "winml": + output_console.print(f"[bold blue]Precision:[/bold blue] {config.precision}") if config.mode != "compare": output_console.print(f"[bold blue]Dataset:[/bold blue] {ds.path}") if ds.name: diff --git a/src/winml/modelkit/eval/fill_mask_evaluator.py b/src/winml/modelkit/eval/fill_mask_evaluator.py index c771f2a6a..65690811f 100644 --- a/src/winml/modelkit/eval/fill_mask_evaluator.py +++ b/src/winml/modelkit/eval/fill_mask_evaluator.py @@ -32,6 +32,8 @@ class WinMLFillMaskEvaluator(WinMLEvaluator): """Evaluate MLMs via pseudo-perplexity.""" + supports_native = False + def __init__( self, config: WinMLEvaluationConfig, @@ -43,7 +45,10 @@ def __init__( mapping = config.dataset.columns_mapping self._input_col = mapping.get("input_column", get_default("fill-mask", "input_column")) - self._tokenizer = AutoTokenizer.from_pretrained(config.model_id) + self._tokenizer = AutoTokenizer.from_pretrained( + config.model_id, + trust_remote_code=config.trust_remote_code, + ) super().__init__(config, model) def prepare_pipeline(self) -> Pipeline: @@ -124,13 +129,13 @@ def compute(self) -> dict[str, Any]: continue encoding = { - k: v for k, v in tok(text, **tok_kwargs).items() - if isinstance(v, torch.Tensor) + k: v for k, v in tok(text, **tok_kwargs).items() if isinstance(v, torch.Tensor) } ids = encoding["input_ids"][0].tolist() specials = tok.get_special_tokens_mask(ids, already_has_special_tokens=True) positions = [ - i for i, (t, s) in enumerate(zip(ids, specials, strict=True)) + i + for i, (t, s) in enumerate(zip(ids, specials, strict=True)) if not s and t != tok.pad_token_id ] if not positions: diff --git a/src/winml/modelkit/eval/keypoint_detection_evaluator.py b/src/winml/modelkit/eval/keypoint_detection_evaluator.py index a7ceeec25..1a7678567 100644 --- a/src/winml/modelkit/eval/keypoint_detection_evaluator.py +++ b/src/winml/modelkit/eval/keypoint_detection_evaluator.py @@ -35,6 +35,8 @@ class WinMLKeypointDetectionEvaluator(WinMLEvaluator): """Evaluator for keypoint detection using COCO OKS-based AP.""" + supports_native = False + def __init__( self, config: WinMLEvaluationConfig, @@ -88,7 +90,10 @@ def prepare_pipeline(self) -> Any: """ from transformers import AutoImageProcessor - processor = AutoImageProcessor.from_pretrained(self.config.model_id) + processor = AutoImageProcessor.from_pretrained( + self.config.model_id, + trust_remote_code=self.config.trust_remote_code, + ) io_config = getattr(self.model, "io_config", None) or {} input_shapes = io_config.get("input_shapes", []) diff --git a/src/winml/modelkit/eval/mask_generation_evaluator.py b/src/winml/modelkit/eval/mask_generation_evaluator.py index 0f5684f9e..e96120ca2 100644 --- a/src/winml/modelkit/eval/mask_generation_evaluator.py +++ b/src/winml/modelkit/eval/mask_generation_evaluator.py @@ -131,6 +131,8 @@ class WinMLMaskGenerationEvaluator(WinMLEvaluator): ``WinMLAutoModel`` composite-registry path. """ + supports_native = False + # Required sub-model role names (must appear as keys in # ``config.model_path`` when it is a dict). _ENCODER_ROLE = "image-encoder" @@ -192,7 +194,7 @@ def __init__( # contract and CodeQL's ``py/missing-call-to-init`` rule. The # base ``prepare_pipeline`` is overridden here to return ``None``, # so it is safe to call from ``WinMLEvaluator.__init__``. - super().__init__(config, model) # type: ignore[arg-type] + super().__init__(config, model) # ------------------------------------------------------------------ # WinMLEvaluator overrides diff --git a/src/winml/modelkit/eval/text_generation_evaluator.py b/src/winml/modelkit/eval/text_generation_evaluator.py index 06026c82a..abc834f43 100644 --- a/src/winml/modelkit/eval/text_generation_evaluator.py +++ b/src/winml/modelkit/eval/text_generation_evaluator.py @@ -54,6 +54,7 @@ class WinMLTextGenerationEvaluator(WinMLEvaluator): * ``seqlen`` -- non-overlapping block length. """ + supports_native = False _TASK = "text-generation" def prepare_pipeline(self) -> Pipeline | None: # type: ignore[override] diff --git a/src/winml/modelkit/eval/zero_shot_classification_evaluator.py b/src/winml/modelkit/eval/zero_shot_classification_evaluator.py index 8e458cc67..46ef68eea 100644 --- a/src/winml/modelkit/eval/zero_shot_classification_evaluator.py +++ b/src/winml/modelkit/eval/zero_shot_classification_evaluator.py @@ -81,12 +81,19 @@ def prepare_pipeline(self) -> Pipeline: # WinMLPreTrainedModel isn't in transformers' Pipeline model union; # the pipeline_class override is also outside the Literal overloads. - pipe = pipeline( # type: ignore[call-overload] - "zero-shot-classification", - model=self.model, - tokenizer=self.config.model_id, - device="cpu", - pipeline_class=_FixedShapeZeroShotPipeline, + pipeline_kwargs: dict[str, Any] = {} + if self.config.trust_remote_code: + pipeline_kwargs["trust_remote_code"] = True + pipe = cast( + "_FixedShapeZeroShotPipeline", + pipeline( + "zero-shot-classification", + model=self.model, + tokenizer=self.config.model_id, + device=self.config.pipeline_device, + pipeline_class=_FixedShapeZeroShotPipeline, + **pipeline_kwargs, + ), ) pipe._winml_evaluator = self diff --git a/src/winml/modelkit/inference/pipeline.py b/src/winml/modelkit/inference/pipeline.py index e14de7af7..36b56ee58 100644 --- a/src/winml/modelkit/inference/pipeline.py +++ b/src/winml/modelkit/inference/pipeline.py @@ -22,19 +22,17 @@ import logging import warnings from collections.abc import Iterable, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any +from typing import Any -if TYPE_CHECKING: - from ..models.winml.base import WinMLPreTrainedModel - from ..models.winml.composite_model import WinMLCompositeModel - logger = logging.getLogger(__name__) # Tasks that WinML recognises but HF ``transformers.pipeline`` does not. # Mapped to their HF pipeline equivalent before calling ``pipeline()``. _HF_PIPELINE_TASK_MAP: dict[str, str] = { "image-to-text": "image-text-to-text", + "next-sentence-prediction": "text-classification", + "sequence-classification": "text-classification", "sentence-similarity": "feature-extraction", } @@ -70,9 +68,10 @@ class _ExtractiveQuestionAnsweringPipeline: task = "question-answering" - def __init__(self, model: Any, tokenizer: Any) -> None: + def __init__(self, model: Any, tokenizer: Any, device: str = "cpu") -> None: self.model = model self.tokenizer = tokenizer + self.device = device self._preprocess_params: dict[str, Any] = {} def preprocess(self, inputs: Any, **kwargs: Any) -> Any: @@ -372,6 +371,7 @@ def _answer( model_inputs[name] = batch if not model_inputs: raise ValueError("Tokenizer outputs do not match the model's declared inputs.") + model_inputs = {name: tensor.to(self.device) for name, tensor in model_inputs.items()} outputs = self.model(**model_inputs) start_logits = getattr(outputs, "start_logits", None) end_logits = getattr(outputs, "end_logits", None) @@ -484,6 +484,8 @@ def _answer( def _create_extractive_question_answering_pipeline( model: Any, model_id: str | None, + device: str = "cpu", + trust_remote_code: bool = False, ) -> _ExtractiveQuestionAnsweringPipeline: """Create the extractive QA pipeline removed from Transformers 5.""" from transformers import AutoTokenizer @@ -491,13 +493,16 @@ def _create_extractive_question_answering_pipeline( tokenizer_source = model_id or getattr(getattr(model, "config", None), "_name_or_path", None) if not tokenizer_source: raise ValueError("model_id is required to load a question-answering tokenizer.") - tokenizer = AutoTokenizer.from_pretrained(tokenizer_source, use_fast=True) + tokenizer_kwargs = {"use_fast": True} + if trust_remote_code: + tokenizer_kwargs["trust_remote_code"] = True + tokenizer = AutoTokenizer.from_pretrained(tokenizer_source, **tokenizer_kwargs) if not getattr(tokenizer, "is_fast", False): raise ValueError( "Extractive question answering requires a fast tokenizer with offset mappings. " "Use a model that provides a fast tokenizer implementation." ) - return _ExtractiveQuestionAnsweringPipeline(model, tokenizer) + return _ExtractiveQuestionAnsweringPipeline(model, tokenizer, device=device) _COMPAT_PIPELINE_FACTORIES = { @@ -523,8 +528,11 @@ def _pipeline_component_kwargs(task: str, model_id: str | None) -> dict[str, str def create_pipeline( task: str, - model: WinMLPreTrainedModel | WinMLCompositeModel, + model: Any, model_id: str | None = None, + *, + device: str = "cpu", + trust_remote_code: bool = False, ) -> Any: """Create an HF pipeline for a WinML model. @@ -533,9 +541,11 @@ def create_pipeline( Args: task: HF task name (e.g. "image-classification") - model: Loaded WinMLPreTrainedModel instance + model: Loaded WinML or native Hugging Face model instance. model_id: HF model ID for loading processors (tokenizer, image processor). If None, pipeline will attempt auto-detection. + device: Device used by the Transformers pipeline for input tensors. + trust_remote_code: Whether custom Hugging Face component code may execute. Returns: A configured task callable ready for inference. @@ -545,14 +555,21 @@ def create_pipeline( hf_task = _HF_PIPELINE_TASK_MAP.get(task, task) compatibility_factory = _COMPAT_PIPELINE_FACTORIES.get(hf_task) if compatibility_factory is not None: - pipe = compatibility_factory(model, model_id) + pipe = compatibility_factory( + model, + model_id, + device=device, + trust_remote_code=trust_remote_code, + ) else: kwargs: dict[str, Any] = { # "device" is for HF pipeline tensor placement, not ORT EP. # WinMLSession handles device delegation internally. - "device": "cpu", + "device": device, **_pipeline_component_kwargs(hf_task, model_id), } + if trust_remote_code: + kwargs["trust_remote_code"] = True # transformers.pipeline has 60+ Literal overloads — runtime task strings can't # be statically matched. The string-task fallback handles unknown tasks safely. diff --git a/src/winml/modelkit/loader/__init__.py b/src/winml/modelkit/loader/__init__.py index a0686b56b..8e1aa16c9 100644 --- a/src/winml/modelkit/loader/__init__.py +++ b/src/winml/modelkit/loader/__init__.py @@ -29,6 +29,7 @@ from ._autoconfig import load_hf_config from .config import WinMLLoaderConfig, resolve_loader_config +from .native import NativeDevice, NativeHFModel, load_native_hf_model, resolve_native_device from .onnx_hub import resolve_hf_onnx_path from .resolution import ( TaskResolution, @@ -53,6 +54,8 @@ "HF_TASK_DEFAULTS", "KNOWN_TASKS", "TASK_SYNONYM_EXTENSIONS", + "NativeDevice", + "NativeHFModel", "TaskResolution", "TaskSource", "WinMLLoaderConfig", @@ -61,11 +64,13 @@ "get_task_abbrev", "load_hf_config", "load_hf_model", + "load_native_hf_model", "normalize_task", "resolve_composite", "resolve_hf_model_class", "resolve_hf_onnx_path", "resolve_loader_config", + "resolve_native_device", "resolve_optimum_library", "resolve_task", "to_optimum_task", diff --git a/src/winml/modelkit/loader/hf.py b/src/winml/modelkit/loader/hf.py index 131dc1179..6f1f7491c 100644 --- a/src/winml/modelkit/loader/hf.py +++ b/src/winml/modelkit/loader/hf.py @@ -145,6 +145,9 @@ def load_hf_model( trust_remote_code: bool = False, hf_config: PretrainedConfig | None = None, model_type: str | None = None, + *, + use_checkpoint_class: bool = False, + torch_dtype: Any | None = None, ) -> tuple[nn.Module, PretrainedConfig, str]: """Load, detect task, and prepare HuggingFace model. @@ -172,6 +175,12 @@ def load_hf_model( hf_config: Optional pre-loaded HF config. When supplied, the ``AutoConfig.from_pretrained`` round-trip is skipped — same dedup pattern as ``resolve_loader_config(hf_config=...)`` from PR #719. + use_checkpoint_class: Load the architecture declared by the checkpoint + instead of a WinML task-specific export wrapper. Falls back to the + task-resolved class when the declared architecture is unavailable + from transformers (for example, a remote-code model). + torch_dtype: Optional dtype policy forwarded to ``from_pretrained``. + Pass ``"auto"`` to preserve the checkpoint's stored dtype. Returns: Tuple of (model, hf_config, task) @@ -263,6 +272,17 @@ def load_hf_model( raise ValueError( f"Cannot resolve task/model for {model_name_or_path}. Original error: {e}" ) from e + if use_checkpoint_class: + from .resolution import _resolve_model_class_from_config + + try: + resolved_class = _resolve_model_class_from_config(hf_config) + except ValueError: + logger.debug( + "Checkpoint architecture is not importable from transformers; " + "using the task-resolved model class %s", + resolved_class.__name__, + ) # [4] Model Instantiation logger.debug("Loading model with class: %s", resolved_class.__name__) @@ -279,11 +299,13 @@ def load_hf_model( if len(matching_subconfigs) == 1: model_config = cast("PretrainedConfig", matching_subconfigs[0]) - model = loader_cls.from_pretrained( - model_name_or_path, - trust_remote_code=trust_remote_code, - config=model_config, - ) + load_kwargs: dict[str, Any] = { + "trust_remote_code": trust_remote_code, + "config": model_config, + } + if torch_dtype is not None: + load_kwargs["torch_dtype"] = torch_dtype + model = loader_cls.from_pretrained(model_name_or_path, **load_kwargs) # [5] Export Preparation model.eval() diff --git a/src/winml/modelkit/loader/native.py b/src/winml/modelkit/loader/native.py new file mode 100644 index 000000000..4f022df4d --- /dev/null +++ b/src/winml/modelkit/loader/native.py @@ -0,0 +1,80 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Native Hugging Face PyTorch loading and device placement.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + + +if TYPE_CHECKING: + from torch import nn + + +@dataclass(frozen=True) +class NativeDevice: + """Resolved PyTorch device and its WinML CLI name.""" + + name: str + torch_device: Any + + +@dataclass(frozen=True) +class NativeHFModel: + """Loaded native Hugging Face model and resolved runtime metadata.""" + + model: "nn.Module" + config: Any + task: str + device: NativeDevice + + +def resolve_native_device(device: str) -> NativeDevice: + """Map a WinML device name to a PyTorch CPU or CUDA device.""" + import torch + + requested = device.lower() + if requested == "auto": + requested = "gpu" if torch.cuda.is_available() else "cpu" + if requested == "gpu": + if not torch.cuda.is_available(): + raise ValueError( + "--device gpu with --runtime pytorch requires a CUDA-enabled PyTorch " + "installation and an available CUDA device." + ) + return NativeDevice(name="gpu", torch_device=torch.device("cuda")) + if requested == "cpu": + return NativeDevice(name="cpu", torch_device=torch.device("cpu")) + raise ValueError( + f"--device {device} is not supported with --runtime pytorch; use auto, cpu, or gpu." + ) + + +def load_native_hf_model( + model_id: str, + *, + task: str | None = None, + device: str = "auto", + trust_remote_code: bool = False, +) -> NativeHFModel: + """Load a checkpoint-declared Hugging Face class without ONNX export.""" + from .hf import load_hf_model + + resolved_device = resolve_native_device(device) + model, hf_config, resolved_task = load_hf_model( + model_id, + task=task, + trust_remote_code=trust_remote_code, + use_checkpoint_class=True, + torch_dtype="auto", + ) + model = model.to(resolved_device.torch_device).eval() + return NativeHFModel( + model=model, + config=hf_config, + task=resolved_task, + device=resolved_device, + ) diff --git a/tests/e2e/test_eval_e2e.py b/tests/e2e/test_eval_e2e.py index b5584c3ba..d73f02e06 100644 --- a/tests/e2e/test_eval_e2e.py +++ b/tests/e2e/test_eval_e2e.py @@ -1093,7 +1093,97 @@ def test_compare_mode_image_classification( # =========================================================================== -# G. CLI-validation error paths (fast — no model load) +# G. Native Hugging Face evaluation +# =========================================================================== + + +class TestEvalNativeHuggingFace: + """Evaluate a real small checkpoint without exporting it to ONNX.""" + + @staticmethod + def _dataset(tmp_path: Path) -> Path: + from datasets import Dataset + + path = tmp_path / "native_text_classification" + Dataset.from_dict( + { + "text": ["A good result.", "A bad result."], + "label": [1, 0], + } + ).save_to_disk(path) + return path + + def test_pytorch_runtime_labeled_evaluation( + self, + runner: CliRunner, + tmp_path: Path, + ) -> None: + out = tmp_path / "native_cpu.json" + _invoke( + runner, + [ + "-m", + "hf-internal-testing/tiny-random-BertForSequenceClassification", + "--task", + "text-classification", + "--dataset", + str(self._dataset(tmp_path)), + "--samples", + "2", + "--no-shuffle", + "--runtime", + "pytorch", + "--device", + "cpu", + "-o", + str(out), + ], + ) + + data = _assert_metrics_present(out, ["accuracy"]) + assert data["runtime"] == "pytorch" + assert data["device"] == "cpu" + assert data["dataset"]["samples"] == 2 + + def test_pytorch_runtime_labeled_evaluation_cuda( + self, + runner: CliRunner, + tmp_path: Path, + ) -> None: + import torch + + if not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + + out = tmp_path / "native_gpu.json" + _invoke( + runner, + [ + "-m", + "hf-internal-testing/tiny-random-BertForSequenceClassification", + "--task", + "text-classification", + "--dataset", + str(self._dataset(tmp_path)), + "--samples", + "2", + "--no-shuffle", + "--runtime", + "pytorch", + "--device", + "gpu", + "-o", + str(out), + ], + ) + + data = _assert_metrics_present(out, ["accuracy"]) + assert data["runtime"] == "pytorch" + assert data["device"] == "gpu" + + +# =========================================================================== +# H. CLI-validation error paths (fast — no model load) # =========================================================================== diff --git a/tests/unit/commands/test_eval_pytorch.py b/tests/unit/commands/test_eval_pytorch.py new file mode 100644 index 000000000..081e5e995 --- /dev/null +++ b/tests/unit/commands/test_eval_pytorch.py @@ -0,0 +1,447 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for native Hugging Face PyTorch evaluation.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +from click.testing import CliRunner + +from winml.modelkit.commands.eval import eval +from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig +from winml.modelkit.loader import NativeDevice, NativeHFModel + + +class TestPyTorchRuntimeCli: + def test_help_shows_runtime_choices(self) -> None: + result = CliRunner().invoke(eval, ["--help"]) + + assert result.exit_code == 0 + assert "--runtime [winml|pytorch]" in result.output + + def test_pytorch_runtime_dispatches_pytorch(self, tmp_path) -> None: + captured: dict[str, WinMLEvaluationConfig] = {} + + def fake_evaluate(config: WinMLEvaluationConfig) -> SimpleNamespace: + captured["config"] = config + return SimpleNamespace(config=config, metrics={}, to_dict=lambda: config.to_dict()) + + with ( + patch("winml.modelkit.eval.evaluate", side_effect=fake_evaluate), + patch("winml.modelkit.commands.eval._write_and_display"), + ): + result = CliRunner().invoke( + eval, + [ + "-m", + "fake/model", + "--task", + "image-classification", + "--dataset", + "fake/dataset", + "--runtime", + "pytorch", + "--device", + "cpu", + "-o", + str(tmp_path / "result.json"), + ], + obj={}, + ) + + assert result.exit_code == 0, result.output + config = captured["config"] + assert config.runtime == "pytorch" + assert config.device == "cpu" + assert config.model_id == "fake/model" + assert config.model_path is None + + def test_default_path_still_exports(self) -> None: + captured: dict[str, WinMLEvaluationConfig] = {} + + def fake_evaluate(config: WinMLEvaluationConfig) -> SimpleNamespace: + captured["config"] = config + return SimpleNamespace(config=config, metrics={}, to_dict=lambda: config.to_dict()) + + with ( + patch("winml.modelkit.eval.evaluate", side_effect=fake_evaluate), + patch("winml.modelkit.commands.eval._resolve_device"), + patch("winml.modelkit.commands.eval._write_and_display"), + ): + result = CliRunner().invoke( + eval, + [ + "-m", + "fake/model", + "--task", + "image-classification", + "--dataset", + "fake/dataset", + ], + obj={}, + ) + + assert result.exit_code == 0, result.output + assert captured["config"].runtime == "winml" + + def test_pytorch_runtime_loads_from_config_file(self, tmp_path) -> None: + config_path = tmp_path / "eval.json" + config_path.write_text( + json.dumps( + { + "eval": { + "runtime": "pytorch", + "model_id": "fake/model", + "task": "image-classification", + "device": "cpu", + "trust_remote_code": True, + "dataset": {"path": "fake/dataset"}, + } + } + ), + encoding="utf-8", + ) + captured: dict[str, WinMLEvaluationConfig] = {} + + def fake_evaluate(config: WinMLEvaluationConfig) -> SimpleNamespace: + captured["config"] = config + return SimpleNamespace(config=config, metrics={}, to_dict=lambda: config.to_dict()) + + with ( + patch("winml.modelkit.eval.evaluate", side_effect=fake_evaluate), + patch("winml.modelkit.commands.eval._write_and_display"), + patch("winml.modelkit.commands.eval._run_dataset_script") as run_dataset_script, + ): + result = CliRunner().invoke(eval, ["--config", str(config_path)], obj={}) + + assert result.exit_code == 0, result.output + assert captured["config"].runtime == "pytorch" + run_dataset_script.assert_called_once_with(captured["config"], True) + + @pytest.mark.parametrize("legacy_field", ["backend", "export_model"]) + def test_config_file_rejects_legacy_runtime_field(self, tmp_path, legacy_field) -> None: + config_path = tmp_path / "eval.json" + config_path.write_text( + json.dumps({"eval": {legacy_field: "pytorch"}}), + encoding="utf-8", + ) + + result = CliRunner().invoke(eval, ["--config", str(config_path)], obj={}) + + assert result.exit_code == 2 + assert "Unsupported eval runtime field" in result.output + assert "Use 'runtime'" in result.output + + @pytest.mark.parametrize( + ("field", "value"), + [ + ("ep", "cpu"), + ("use_cache", True), + ("skip_build", True), + ("quant", True), + ], + ) + def test_native_config_file_rejects_onnx_only_fields( + self, + tmp_path, + field: str, + value: object, + ) -> None: + config_path = tmp_path / "eval.json" + config_path.write_text( + json.dumps( + { + "eval": { + "runtime": "pytorch", + "model_id": "fake/model", + "task": "image-classification", + field: value, + } + } + ), + encoding="utf-8", + ) + + result = CliRunner().invoke( + eval, + ["-m", "fake/model", "--config", str(config_path)], + obj={}, + ) + + assert result.exit_code == 2 + assert f"eval.{field}" in result.output + + @pytest.mark.parametrize( + ("args", "expected_flag"), + [ + (["--ep", "cpu"], "--ep"), + (["--precision", "fp16"], "--precision"), + (["--no-quant"], "--quant/--no-quant"), + (["--no-optimize"], "--optimize/--no-optimize"), + (["--no-analyze"], "--analyze/--no-analyze"), + (["--max-optim-iterations", "2"], "--max-optim-iterations"), + (["--allow-unsupported-nodes"], "--allow-unsupported-nodes"), + (["--no-skip-build"], "--skip-build/--no-skip-build"), + (["--no-use-cache"], "--use-cache/--no-use-cache"), + (["--rebuild"], "--rebuild/--no-rebuild"), + (["--mode", "compare"], "--mode"), + ], + ) + def test_rejects_onnx_only_options( + self, + args: list[str], + expected_flag: str, + ) -> None: + result = CliRunner().invoke( + eval, + ["-m", "fake/model", "--runtime", "pytorch", *args], + obj={}, + ) + + assert result.exit_code == 2 + assert expected_flag in result.output + + def test_rejects_export_override(self, tmp_path) -> None: + shape_config = tmp_path / "shape.json" + shape_config.write_text(json.dumps({"height": 16})) + + result = CliRunner().invoke( + eval, + [ + "-m", + "fake/model", + "--runtime", + "pytorch", + "--shape-config", + str(shape_config), + ], + obj={}, + ) + + assert result.exit_code == 2 + assert "--shape-config" in result.output + + def test_rejects_onnx_input(self, tmp_path) -> None: + model_path = tmp_path / "model.onnx" + model_path.write_bytes(b"not used") + + result = CliRunner().invoke( + eval, + [ + "-m", + str(model_path), + "--model-id", + "fake/model", + "--runtime", + "pytorch", + ], + obj={}, + ) + + assert result.exit_code == 2 + assert "requires a Hugging Face model ID" in result.output + + def test_rejects_genai_bundle(self, tmp_path) -> None: + (tmp_path / "genai_config.json").write_text("{}") + + result = CliRunner().invoke( + eval, + ["-m", str(tmp_path), "--runtime", "pytorch"], + obj={}, + ) + + assert result.exit_code == 2 + assert "GenAI bundles are not supported" in result.output + + def test_rejects_npu_device(self) -> None: + result = CliRunner().invoke( + eval, + ["-m", "fake/model", "--runtime", "pytorch", "--device", "npu"], + obj={}, + ) + + assert result.exit_code == 2 + assert "use auto, cpu, or gpu" in result.output + + @pytest.mark.parametrize( + "task", + [ + "fill-mask", + "keypoint-detection", + "mask-generation", + "text-generation", + ], + ) + def test_rejects_non_pipeline_evaluators(self, task: str) -> None: + result = CliRunner().invoke( + eval, + ["-m", "fake/model", "--runtime", "pytorch", "--task", task], + obj={}, + ) + + assert result.exit_code == 2 + assert "does not use the standard labeled Hugging Face pipeline" in result.output + + def test_gpu_requires_cuda(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + result = CliRunner().invoke( + eval, + ["-m", "fake/model", "--runtime", "pytorch", "--device", "gpu"], + obj={}, + ) + + assert result.exit_code == 2 + assert "requires a CUDA-enabled PyTorch" in result.output + + +class TestNativeEvaluation: + def test_pytorch_runtime_uses_pytorch_loader_kind(self) -> None: + from winml.modelkit.eval.evaluate import _ModelLoaderKind, _select_model_loader + + config = WinMLEvaluationConfig( + model_id="fake/model", + task="image-classification", + runtime="pytorch", + ) + + assert _select_model_loader(config) is _ModelLoaderKind.PYTORCH + + def test_public_evaluate_rejects_invalid_runtime(self) -> None: + from typing import cast + + from winml.modelkit.eval import EvalRuntime, evaluate + + config = WinMLEvaluationConfig( + model_id="fake/model", + task="image-classification", + runtime=cast("EvalRuntime", "invalid"), + ) + + with pytest.raises(ValueError, match="Invalid runtime"): + evaluate(config) + + def test_public_evaluate_rejects_onnx_state(self) -> None: + from winml.modelkit.eval import evaluate + + config = WinMLEvaluationConfig( + model_id="fake/model", + model_path="model.onnx", + task="image-classification", + runtime="pytorch", + ) + + with pytest.raises(ValueError, match="model_path"): + evaluate(config) + + @pytest.mark.parametrize( + ("config_override", "expected_field"), + [ + ({"use_cache": False}, "use_cache"), + ({"rebuild": True}, "rebuild"), + ], + ) + def test_public_evaluate_rejects_cache_state( + self, + config_override: dict[str, bool], + expected_field: str, + ) -> None: + from winml.modelkit.eval import evaluate + + config = WinMLEvaluationConfig( + model_id="fake/model", + task="image-classification", + runtime="pytorch", + **config_override, + ) + + with pytest.raises(ValueError, match=expected_field): + evaluate(config) + + def test_load_model_uses_shared_native_loader(self) -> None: + from winml.modelkit.eval.evaluate import _load_model + + model = MagicMock() + loaded = NativeHFModel( + model=model, + config=MagicMock(), + task="image-classification", + device=NativeDevice(name="gpu", torch_device=torch.device("cuda")), + ) + config = WinMLEvaluationConfig( + model_id="fake/model", + task="image-classification", + device="gpu", + runtime="pytorch", + trust_remote_code=True, + ) + + with patch( + "winml.modelkit.loader.load_native_hf_model", + return_value=loaded, + ) as load: + assert _load_model(config) is model + + load.assert_called_once_with( + "fake/model", + task="image-classification", + device="gpu", + trust_remote_code=True, + ) + assert config.device == "gpu" + + def test_representative_evaluator_uses_native_pipeline_device(self) -> None: + from winml.modelkit.eval.base_evaluator import WinMLEvaluator + + evaluator = WinMLEvaluator.__new__(WinMLEvaluator) + evaluator.config = WinMLEvaluationConfig( + model_id="fake/model", + task="image-classification", + device="gpu", + runtime="pytorch", + dataset=DatasetConfig(path="fake/dataset"), + ) + evaluator.model = MagicMock() + pipeline = MagicMock() + + with patch( + "winml.modelkit.inference.pipeline.create_pipeline", + return_value=pipeline, + ) as create: + assert evaluator.prepare_pipeline() is pipeline + + create.assert_called_once_with( + "image-classification", + evaluator.model, + "fake/model", + device="cuda", + trust_remote_code=False, + ) + + def test_config_roundtrip_identifies_pytorch_runtime(self) -> None: + config = WinMLEvaluationConfig( + model_id="fake/model", + device="cpu", + runtime="pytorch", + ) + + serialized = config.to_dict() + restored = WinMLEvaluationConfig.from_dict(serialized) + + assert serialized["runtime"] == "pytorch" + assert "skip_build" not in serialized + assert "use_cache" not in serialized + assert "rebuild" not in serialized + assert restored.runtime == "pytorch" + + @pytest.mark.parametrize("legacy_field", ["backend", "export_model"]) + def test_config_deserialization_rejects_legacy_runtime_field(self, legacy_field) -> None: + with pytest.raises(ValueError, match="Use 'runtime'"): + WinMLEvaluationConfig.from_dict({legacy_field: "pytorch"}) diff --git a/tests/unit/eval/test_eval.py b/tests/unit/eval/test_eval.py index fb5babe8c..e41ad6e8e 100644 --- a/tests/unit/eval/test_eval.py +++ b/tests/unit/eval/test_eval.py @@ -28,7 +28,10 @@ def test_relies_on_model_framework_inference(self) -> None: evaluator.model = MagicMock() sentinel = MagicMock() - with patch("transformers.pipeline", return_value=sentinel) as mock_pipeline: + with patch( + "winml.modelkit.inference.pipeline.create_pipeline", + return_value=sentinel, + ) as mock_pipeline: assert evaluator.prepare_pipeline() is sentinel assert "framework" not in mock_pipeline.call_args.kwargs @@ -217,6 +220,28 @@ def test_infer_from_model_id(self): ): assert _resolve_task(config) == "image-classification" + def test_infer_threads_trust_remote_code(self): + from winml.modelkit.eval.evaluate import _resolve_task + + fake_resolution = MagicMock(task="image-classification") + config = WinMLEvaluationConfig( + model_id="custom/model", + trust_remote_code=True, + ) + with ( + patch( + "winml.modelkit.loader.load_hf_config", + return_value=MagicMock(), + ) as load_config, + patch( + "winml.modelkit.loader.resolution.resolve_task", + return_value=fake_resolution, + ), + ): + assert _resolve_task(config) == "image-classification" + + assert load_config.call_args.kwargs["trust_remote_code"] is True + def test_explicit_feature_extraction_preserved_verbatim(self): """Explicit --task is surfaced verbatim (explicit means explicit). @@ -1649,6 +1674,7 @@ def test_load_model_from_onnx(self): model_path="model.onnx", task="image-classification", device="cpu", + trust_remote_code=True, ) with ( @@ -1659,10 +1685,11 @@ def test_load_model_from_onnx(self): patch( "winml.modelkit.loader.load_hf_config", return_value=mock_hf_config, - ), + ) as load_hf_config, ): result = eval_mod._load_model(config) + assert load_hf_config.call_args.kwargs["trust_remote_code"] is True mock_auto.from_onnx.assert_called_once() assert mock_auto.from_onnx.call_args.kwargs["use_cache"] is True assert mock_auto.from_onnx.call_args.kwargs["force_rebuild"] is False diff --git a/tests/unit/eval/test_fill_mask_evaluator.py b/tests/unit/eval/test_fill_mask_evaluator.py index bc5766d1e..514551f38 100644 --- a/tests/unit/eval/test_fill_mask_evaluator.py +++ b/tests/unit/eval/test_fill_mask_evaluator.py @@ -55,8 +55,10 @@ def _make_evaluator(model=None, max_length=None): ), ) - with patch("datasets.load_dataset", return_value=mock_ds), \ - patch("transformers.AutoTokenizer.from_pretrained", return_value=_make_tokenizer()): + with ( + patch("datasets.load_dataset", return_value=mock_ds), + patch("transformers.AutoTokenizer.from_pretrained", return_value=_make_tokenizer()), + ): return WinMLFillMaskEvaluator(config, model) @@ -68,6 +70,37 @@ def test_tokenizer_loaded_in_init(self) -> None: evaluator = _make_evaluator() assert evaluator._tokenizer.mask_token_id == 103 + def test_tokenizer_forwards_trust_remote_code(self) -> None: + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + config = WinMLEvaluationConfig( + model_id="test/mock-bert", + task="fill-mask", + trust_remote_code=True, + dataset=DatasetConfig(path="Salesforce/wikitext"), + ) + model = MagicMock() + model.config.label2id = None + model.io_config = {} + mock_ds = MagicMock() + mock_ds.__len__ = lambda self: 5 + mock_ds.shuffle.return_value = mock_ds + mock_ds.select.return_value = mock_ds + + with ( + patch("datasets.load_dataset", return_value=mock_ds), + patch( + "transformers.AutoTokenizer.from_pretrained", + return_value=_make_tokenizer(), + ) as load_tokenizer, + ): + WinMLFillMaskEvaluator(config, model) + + load_tokenizer.assert_called_once_with( + "test/mock-bert", + trust_remote_code=True, + ) + def test_align_labels_is_noop(self) -> None: evaluator = _make_evaluator() ds = MagicMock() @@ -95,7 +128,8 @@ class TestLogits: def test_dict_with_logits_key(self) -> None: logits = torch.randn(1, 5, 50) assert torch.equal( - _make_evaluator()._logits({"logits": logits, "aux": None}), logits, + _make_evaluator()._logits({"logits": logits, "aux": None}), + logits, ) def test_dict_without_logits_key_raises(self) -> None: diff --git a/tests/unit/eval/test_keypoint_detection_evaluator.py b/tests/unit/eval/test_keypoint_detection_evaluator.py index bfa8217aa..4c03a12c6 100644 --- a/tests/unit/eval/test_keypoint_detection_evaluator.py +++ b/tests/unit/eval/test_keypoint_detection_evaluator.py @@ -12,6 +12,8 @@ from __future__ import annotations +from unittest.mock import MagicMock, patch + import pytest import torch @@ -42,6 +44,29 @@ def test_xyxy_converted_to_xywh(self): ev = _make_evaluator("xyxy") assert ev._to_xywh([10.0, 20.0, 40.0, 60.0]) == [10.0, 20.0, 30.0, 40.0] + def test_processor_forwards_trust_remote_code(self): + from winml.modelkit.eval import WinMLEvaluationConfig + + ev = _make_evaluator() + ev.config = WinMLEvaluationConfig( + model_id="test/model", + task="keypoint-detection", + trust_remote_code=True, + ) + ev.model = MagicMock(io_config={}) + processor = MagicMock() + + with patch( + "transformers.AutoImageProcessor.from_pretrained", + return_value=processor, + ) as load_processor: + assert ev.prepare_pipeline() is processor + + load_processor.assert_called_once_with( + "test/model", + trust_remote_code=True, + ) + class TestPredictionFlattening: def test_flatten_interleaves_xy_and_score(self): @@ -189,4 +214,3 @@ def test_dataset_index_detected_from_forward_signature(self): ev.model = model ev._predict_poses(_MockProcessor(), object(), [[0.0, 0.0, 10.0, 10.0]]) assert model.calls[0]["dataset_index"].tolist() == [2] - diff --git a/tests/unit/inference/test_pipeline.py b/tests/unit/inference/test_pipeline.py index 7755ecdb6..2482e3218 100644 --- a/tests/unit/inference/test_pipeline.py +++ b/tests/unit/inference/test_pipeline.py @@ -160,6 +160,13 @@ def test_sentence_similarity_maps_to_feature_extraction(self) -> None: def test_image_to_text_maps_to_transformers_5_name(self) -> None: assert _HF_PIPELINE_TASK_MAP["image-to-text"] == "image-text-to-text" + @pytest.mark.parametrize( + "task", + ["sequence-classification", "next-sentence-prediction"], + ) + def test_classification_aliases_map_to_transformers_task(self, task: str) -> None: + assert _HF_PIPELINE_TASK_MAP[task] == "text-classification" + def test_unknown_task_not_in_map(self) -> None: assert "image-classification" not in _HF_PIPELINE_TASK_MAP @@ -182,6 +189,55 @@ class ProcessorOnlyPipeline: class TestCreatePipeline: + def test_threads_trust_remote_code_to_pipeline(self) -> None: + model = MagicMock() + pipe = MagicMock() + pipe.tokenizer = None + pipe.image_processor = None + + with ( + patch( + "winml.modelkit.inference.pipeline._pipeline_component_kwargs", + return_value={}, + ), + patch("transformers.pipeline", return_value=pipe) as pipeline, + ): + create_pipeline( + "image-classification", + model, + "test-model", + trust_remote_code=True, + ) + + assert pipeline.call_args.kwargs["trust_remote_code"] is True + + def test_places_native_pipeline_inputs_on_cuda(self) -> None: + model = MagicMock() + pipe = MagicMock() + pipe.tokenizer = None + pipe.image_processor = None + + with ( + patch( + "winml.modelkit.inference.pipeline._pipeline_component_kwargs", + return_value={}, + ), + patch("transformers.pipeline", return_value=pipe) as pipeline, + ): + result = create_pipeline( + "image-classification", + model, + "test-model", + device="cuda", + ) + + assert result is pipe + pipeline.assert_called_once_with( + "image-classification", + model=model, + device="cuda", + ) + def test_constructs_real_transformers_pipeline(self) -> None: from transformers import ViTConfig, ViTForImageClassification, ViTImageProcessor from transformers.pipelines import ImageClassificationPipeline diff --git a/tests/unit/loader/test_load_hf_model.py b/tests/unit/loader/test_load_hf_model.py index 1c3f4066f..cbe0a4766 100644 --- a/tests/unit/loader/test_load_hf_model.py +++ b/tests/unit/loader/test_load_hf_model.py @@ -195,6 +195,52 @@ def mock_resolve(config, *, task=None, model_class=None, model_type_override=Non assert call["task"] is None # Auto-detect assert call["model_class"] is None + def test_checkpoint_class_and_dtype_override_export_wrapper(self, monkeypatch): + """Native consumers can load the checkpoint architecture and stored dtype.""" + from types import SimpleNamespace + from unittest.mock import MagicMock + + import winml.modelkit.loader.resolution as resolution_module + + checkpoint_class = MagicMock() + checkpoint_class.__name__ = "CheckpointModel" + checkpoint_class.config_class = None + checkpoint_model = MagicMock() + checkpoint_class.from_pretrained.return_value = checkpoint_model + export_wrapper = MagicMock() + export_wrapper.__name__ = "ExportWrapper" + config = SimpleNamespace(model_type="unit", architectures=["CheckpointModel"]) + + monkeypatch.setattr( + resolution_module, + "resolve_task", + lambda *_a, **_kw: SimpleNamespace( + task="text-generation", + model_class=export_wrapper, + ), + ) + monkeypatch.setattr( + resolution_module, + "_resolve_model_class_from_config", + lambda _config: checkpoint_class, + ) + + model, _, task = load_hf_model( + "fake/model", + hf_config=config, + use_checkpoint_class=True, + torch_dtype="auto", + ) + + assert model is checkpoint_model + assert task == "text-generation" + checkpoint_class.from_pretrained.assert_called_once_with( + "fake/model", + trust_remote_code=False, + config=config, + torch_dtype="auto", + ) + def test_bert_tiny_uses_model_specific_default_task(self, monkeypatch): """bert-tiny should use model-specific default task when task is omitted.""" from unittest.mock import MagicMock diff --git a/tests/unit/loader/test_native_hf.py b/tests/unit/loader/test_native_hf.py new file mode 100644 index 000000000..112b9f522 --- /dev/null +++ b/tests/unit/loader/test_native_hf.py @@ -0,0 +1,70 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for shared native Hugging Face loading.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from winml.modelkit.loader import load_native_hf_model, resolve_native_device + + +class TestResolveNativeDevice: + def test_auto_uses_cpu_without_cuda(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + resolved = resolve_native_device("auto") + + assert resolved.name == "cpu" + assert resolved.torch_device == torch.device("cpu") + + def test_auto_uses_cuda_when_available(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + + resolved = resolve_native_device("auto") + + assert resolved.name == "gpu" + assert resolved.torch_device == torch.device("cuda") + + def test_gpu_requires_cuda(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + with pytest.raises(ValueError, match="requires a CUDA-enabled PyTorch"): + resolve_native_device("gpu") + + +def test_load_native_hf_model_preserves_checkpoint_class_and_dtype() -> None: + model = MagicMock() + model.to.return_value = model + model.eval.return_value = model + hf_config = MagicMock() + + with patch( + "winml.modelkit.loader.hf.load_hf_model", + return_value=(model, hf_config, "image-classification"), + ) as load: + result = load_native_hf_model( + "fake/model", + task="image-classification", + device="cpu", + trust_remote_code=True, + ) + + load.assert_called_once_with( + "fake/model", + task="image-classification", + trust_remote_code=True, + use_checkpoint_class=True, + torch_dtype="auto", + ) + model.to.assert_called_once_with(torch.device("cpu")) + model.eval.assert_called_once_with() + assert result.model is model + assert result.config is hf_config + assert result.task == "image-classification" + assert result.device.name == "cpu"