support(transformer): add DigitalEdu tabular recipe

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
+}

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
+}

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(

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",

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,
     )

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",
+]

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",

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)