diff --git a/examples/recipes/SnowFlash383935_DigitalEduTransformers/cpu/cpu/tabular-classification_fp16_config.json b/examples/recipes/SnowFlash383935_DigitalEduTransformers/cpu/cpu/tabular-classification_fp16_config.json new file mode 100644 index 000000000..e5fb2d335 --- /dev/null +++ b/examples/recipes/SnowFlash383935_DigitalEduTransformers/cpu/cpu/tabular-classification_fp16_config.json @@ -0,0 +1,30 @@ +{ + "loader": { + "model_type": "transformer", + "model_class": "TabularTransformerWrapper", + "task": "tabular-classification", + "trust_remote_code": true + }, + "export": { + "opset_version": 17, + "batch_size": 1, + "input_tensors": [ + { + "name": "features", + "dtype": "float32", + "shape": [1, 7], + "value_range": [0, 1000] + } + ], + "output_tensors": [ + { + "name": "logits" + } + ] + }, + "optim": {}, + "quant": { + "mode": "fp16" + }, + "compile": null +} diff --git a/examples/recipes/SnowFlash383935_DigitalEduTransformers/cpu/cpu/tabular-classification_fp32_config.json b/examples/recipes/SnowFlash383935_DigitalEduTransformers/cpu/cpu/tabular-classification_fp32_config.json new file mode 100644 index 000000000..e7d92f5bc --- /dev/null +++ b/examples/recipes/SnowFlash383935_DigitalEduTransformers/cpu/cpu/tabular-classification_fp32_config.json @@ -0,0 +1,28 @@ +{ + "loader": { + "model_type": "transformer", + "model_class": "TabularTransformerWrapper", + "task": "tabular-classification", + "trust_remote_code": true + }, + "export": { + "opset_version": 17, + "batch_size": 1, + "input_tensors": [ + { + "name": "features", + "dtype": "float32", + "shape": [1, 7], + "value_range": [0, 1000] + } + ], + "output_tensors": [ + { + "name": "logits" + } + ] + }, + "optim": {}, + "quant": null, + "compile": null +} diff --git a/src/winml/modelkit/inference/tasks.py b/src/winml/modelkit/inference/tasks.py index e0602ff65..99ef2118d 100644 --- a/src/winml/modelkit/inference/tasks.py +++ b/src/winml/modelkit/inference/tasks.py @@ -294,6 +294,17 @@ def _postprocess_sentence_similarity( ], mapping=PipelineMapping(pipe_input="text"), ), + "tabular-classification": TaskInputSpec( + user_inputs=[ + InputField( + name="features", + type="json", + required=True, + description="Numeric feature vector", + ), + ], + mapping=PipelineMapping(pipe_input="features"), + ), "token-classification": TaskInputSpec( user_inputs=[ InputField( diff --git a/src/winml/modelkit/loader/task.py b/src/winml/modelkit/loader/task.py index 0ae0ca6be..1864f4c4c 100644 --- a/src/winml/modelkit/loader/task.py +++ b/src/winml/modelkit/loader/task.py @@ -76,6 +76,7 @@ "next-sentence-prediction": "nsp", "table-question-answering": "tabqa", "document-question-answering": "docqa", + "tabular-classification": "tabcls", "sentence-similarity": None, # Audio "audio-classification": "audiocls", @@ -250,6 +251,9 @@ def normalize_task(task: str) -> str: # collapse to feature-extraction is guarded by tests/unit/loader/test_task_boundary.py. # next-sentence-prediction has the same I/O as text-classification: input_ids -> logits "next-sentence-prediction": "text-classification", + # tabular-classification uses a custom WinML task name but reuses Optimum's + # classification exporter boundary once the model exposes tensor features. + "tabular-classification": "text-classification", # mask-generation is registered via register_onnx_overwrite for SAM2. # Optimum incorrectly maps it to "feature-extraction"; preserve as-is. "mask-generation": "mask-generation", diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index b11e7b416..94cf2f1a5 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/src/winml/modelkit/models/hf/__init__.py @@ -95,6 +95,8 @@ from .t5 import T5_CONFIG from .t5 import T5DecoderIOConfig as _T5DecoderIOConfig # triggers registration from .t5 import T5EncoderIOConfig as _T5EncoderIOConfig # triggers registration +from .transformer import MODEL_CLASS_MAPPING as _TRANSFORMER_CLASS_MAPPING +from .transformer import TabularTransformerIOConfig as _TabularTransformerIOConfig from .vision_encoder_decoder import MODEL_CLASS_MAPPING as _VED_CLASS_MAPPING from .vision_encoder_decoder import VISION_ENCODER_DECODER_CONFIG from .vision_encoder_decoder import ( @@ -133,6 +135,7 @@ _SEGFORMER_CLASS_MAPPING, _SIGLIP_CLASS_MAPPING, _T5_CLASS_MAPPING, + _TRANSFORMER_CLASS_MAPPING, _VED_CLASS_MAPPING, _VITPOSE_CLASS_MAPPING, ) diff --git a/src/winml/modelkit/models/hf/transformer.py b/src/winml/modelkit/models/hf/transformer.py new file mode 100644 index 000000000..1c43a731e --- /dev/null +++ b/src/winml/modelkit/models/hf/transformer.py @@ -0,0 +1,132 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Custom tabular Transformer export support. + +The SnowFlash383935/DigitalEduTransformers checkpoint publishes a custom +``model_type='transformer'`` architecture whose public ``forward`` accepts a +Python list, performs NumPy normalization, and returns Python booleans. That is +useful for the model card example but is not an ONNX export contract. The +wrapper below exposes the tensor path directly: normalized tabular features in, +logits out. +""" + +from __future__ import annotations + +from typing import Any, cast + +import torch +import torch.nn as nn +from optimum.exporters.onnx import OnnxConfig +from optimum.utils import NormalizedConfig +from optimum.utils.input_generators import DummyInputGenerator + +from ...export import register_onnx_overwrite + + +class TabularInputGenerator(DummyInputGenerator): # type: ignore[misc] + """Generate floating tabular feature tensors.""" + + SUPPORTED_INPUT_NAMES = ("features",) + + def __init__( + self, + task: str, + normalized_config: NormalizedConfig, + batch_size: int = 1, + **_: Any, + ) -> None: + """Initialize the generator from Optimum's OnnxConfig factory args.""" + self.task = task + self.normalized_config = normalized_config + self.batch_size = batch_size + + def generate( + self, + input_name: str, + framework: str = "pt", + int_dtype: str = "int64", + float_dtype: str = "fp32", + ) -> Any: + """Generate a zero-valued tabular feature tensor.""" + import torch + + if input_name != "features": + raise ValueError(f"Unsupported input for tabular transformer: {input_name}") + input_dim = int(getattr(self.normalized_config, "input_dim", 7)) + return torch.zeros((self.batch_size, input_dim), dtype=torch.float32) + + +class TabularTransformerWrapper(nn.Module): + """Tensor-in/tensor-out wrapper around the remote tabular model.""" + + def __init__(self, model: nn.Module) -> None: + super().__init__() + self.config = model.config + self.input_proj = cast("nn.Module", model.input_proj) + self.transformer = cast("nn.Module", model.transformer) + self.head = cast("nn.Module", model.head) + mean = torch.tensor(self.config.mean, dtype=torch.float32) + std = torch.tensor(self.config.std, dtype=torch.float32) + self.register_buffer("mean", mean) + self.register_buffer("std", std) + + @classmethod + def from_pretrained(cls, model_name_or_path: str, **kwargs: Any) -> TabularTransformerWrapper: + """Load the remote model and wrap its tensorizable submodules.""" + from transformers import AutoModel + + model = AutoModel.from_pretrained(model_name_or_path, **kwargs) + wrapper = cls(model) + wrapper.eval() + return wrapper + + def forward(self, features: torch.Tensor) -> torch.Tensor: + """Run normalized tabular features through the Transformer classifier.""" + mean = cast("torch.Tensor", self.mean) + std = cast("torch.Tensor", self.std) + x = (features - mean.to(features.device, features.dtype)) / ( + std.to(features.device, features.dtype) + 1e-8 + ) + x = self.input_proj(x) + x = x.unsqueeze(1) + x = self.transformer(x) + x = x.squeeze(1) + return cast("torch.Tensor", self.head(x)) + + +@register_onnx_overwrite("transformer", "text-classification", library_name="transformers") +class TabularTransformerIOConfig(OnnxConfig): # type: ignore[misc] + """ONNX config for tensorized tabular classification.""" + + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + input_dim="input_dim", + allow_new=True, + ) + DUMMY_INPUT_GENERATOR_CLASSES = (TabularInputGenerator,) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + """ONNX input names and dynamic axes.""" + return {"features": {0: "batch_size"}} + + @property + def outputs(self) -> dict[str, dict[int, str]]: + """ONNX output names and dynamic axes.""" + return {"logits": {0: "batch_size"}} + + +MODEL_CLASS_MAPPING: dict[tuple[str, str | None], type] = { + ("transformer", "tabular-classification"): TabularTransformerWrapper, + ("transformer", "text-classification"): TabularTransformerWrapper, + ("transformer", None): TabularTransformerWrapper, +} + + +__all__ = [ + "MODEL_CLASS_MAPPING", + "TabularInputGenerator", + "TabularTransformerIOConfig", + "TabularTransformerWrapper", +] diff --git a/src/winml/modelkit/models/winml/__init__.py b/src/winml/modelkit/models/winml/__init__.py index bd1ac303b..0fc307e5b 100644 --- a/src/winml/modelkit/models/winml/__init__.py +++ b/src/winml/modelkit/models/winml/__init__.py @@ -37,6 +37,7 @@ "text-classification": "WinMLModelForSequenceClassification", "sequence-classification": "WinMLModelForSequenceClassification", "next-sentence-prediction": "WinMLModelForSequenceClassification", + "tabular-classification": "WinMLModelForSequenceClassification", "image-segmentation": "WinMLModelForImageSegmentation", "semantic-segmentation": "WinMLModelForSemanticSegmentation", "object-detection": "WinMLModelForObjectDetection", diff --git a/tests/unit/models/test_tabular_transformer.py b/tests/unit/models/test_tabular_transformer.py new file mode 100644 index 000000000..1d45dff12 --- /dev/null +++ b/tests/unit/models/test_tabular_transformer.py @@ -0,0 +1,55 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from __future__ import annotations + +from typing import ClassVar + +import torch + +from winml.modelkit.loader.task import TASK_SYNONYM_EXTENSIONS, to_optimum_task +from winml.modelkit.models.hf import MODEL_CLASS_MAPPING +from winml.modelkit.models.hf.transformer import TabularTransformerWrapper + + +class _Config: + mean: ClassVar[list[float]] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0] + std: ClassVar[list[float]] = [1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0] + + +class _RemoteTabularModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = _Config() + self.input_proj = torch.nn.Linear(7, 3, bias=False) + self.transformer = torch.nn.Identity() + self.head = torch.nn.Linear(3, 1, bias=False) + + +def test_tabular_task_maps_to_optimum_text_classification() -> None: + assert TASK_SYNONYM_EXTENSIONS["tabular-classification"] == "text-classification" + assert to_optimum_task("tabular-classification") == "text-classification" + + +def test_transformer_model_class_mapping_registers_tabular_wrapper() -> None: + assert ( + MODEL_CLASS_MAPPING[("transformer", "tabular-classification")] + is TabularTransformerWrapper + ) + assert MODEL_CLASS_MAPPING[("transformer", "text-classification")] is TabularTransformerWrapper + assert MODEL_CLASS_MAPPING[("transformer", None)] is TabularTransformerWrapper + + +def test_tabular_wrapper_exposes_tensor_logits() -> None: + remote = _RemoteTabularModel() + wrapper = TabularTransformerWrapper(remote) + + features = torch.tensor([[1.0, 4.0, 7.0, 12.0, 21.0, 38.0, 71.0]]) + logits = wrapper(features) + + normalized = (features - torch.tensor(_Config.mean)) / (torch.tensor(_Config.std) + 1e-8) + hidden = remote.input_proj(normalized).unsqueeze(1) + expected = remote.head(remote.transformer(hidden).squeeze(1)) + torch.testing.assert_close(logits, expected) + assert logits.shape == (1, 1)