feat: export classifier and regressor pipelines to ONNX

scripts/export_to_onnx.py

@@ -17,6 +17,8 @@
 from transformers import AutoTokenizer, PreTrainedTokenizerFast
 
 from model2vec import StaticModel
+from model2vec.inference import StaticModelPipeline
+from model2vec.inference.mlp import Activation
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -80,8 +82,63 @@ def tokenize(self, sentences: list[str], max_length: int | None = None) -> tuple
         return input_ids, offsets
 
 
+class TorchStaticModelPipeline(torch.nn.Module):
+    def __init__(self, pipeline: StaticModelPipeline) -> None:
+        """Wrap a StaticModelPipeline (encoder + MLP head) as a single torch module."""
+        super().__init__()
+        self.encoder = TorchStaticModel(pipeline.model)
+        self.activation = pipeline.head.activation
+        # Rebuild the head Layers as nn.Linear. Layer stores weight as [out, in] and computes
+        # x @ weight.T + bias, which matches nn.Linear(in, out) exactly.
+        self.layers = torch.nn.ModuleList()
+        for layer in pipeline.head.layers:
+            weight = torch.from_numpy(layer.weight)
+            linear = torch.nn.Linear(weight.shape[1], weight.shape[0])
+            linear.weight = torch.nn.Parameter(weight, requires_grad=False)
+            linear.bias = torch.nn.Parameter(torch.from_numpy(layer.bias), requires_grad=False)
+            self.layers.append(linear)
+
+    def forward(self, input_ids: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
+        """Encode the inputs and run them through the head, applying the output activation."""
+        out = self.encoder(input_ids, offsets)
+        *hidden_layers, last_layer = self.layers
+        for layer in hidden_layers:
+            out = torch.relu(layer(out))
+        logits = last_layer(out)
+        if self.activation == Activation.SOFTMAX:
+            return torch.softmax(logits, dim=-1)
+        if self.activation == Activation.SIGMOID:
+            return torch.sigmoid(logits)
+        return logits
+
+    def tokenize(self, sentences: list[str], max_length: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
+        """Tokenize the input sentences using the underlying encoder."""
+        return self.encoder.tokenize(sentences, max_length)
+
+
 def export_model_to_onnx(model_path: str, save_path: Path) -> None:
-    """Export the StaticModel to ONNX format and save tokenizer files.
+    """Export a StaticModel or a StaticModelPipeline to ONNX format and save tokenizer files.
+
+    A classifier/regressor pipeline (one with a trained head) is detected automatically and
+    exported with its head included; a plain encoder is exported as embeddings.
+
+    :param model_path: The path to the pretrained StaticModel or StaticModelPipeline.
+    :param save_path: The directory to save the model and related files.
+    """
+    try:
+        pipeline = StaticModelPipeline.from_pretrained(model_path)
+    except FileNotFoundError:
+        pipeline = None
+
+    if pipeline is not None:
+        export_pipeline_to_onnx(pipeline, save_path)
+        return
+
+    _export_encoder_to_onnx(model_path, save_path)
+
+
+def _export_encoder_to_onnx(model_path: str, save_path: Path) -> None:
+    """Export a plain StaticModel encoder to ONNX format and save tokenizer files.
 
     :param model_path: The path to the pretrained StaticModel.
     :param save_path: The directory to save the model and related files.
@@ -125,6 +182,55 @@ def export_model_to_onnx(model_path: str, save_path: Path) -> None:
     logger.info(f"Tokenizer files have been saved to {save_path}")
 
 
+def export_pipeline_to_onnx(pipeline: StaticModelPipeline, save_path: Path) -> None:
+    """Export a StaticModelPipeline (encoder + classifier/regressor head) to ONNX format.
+
+    The exported graph outputs class probabilities for classifiers (softmax/sigmoid heads) or
+    raw predictions for regression/projector heads (identity activation).
+
+    :param pipeline: The pretrained StaticModelPipeline.
+    :param save_path: The directory to save the model and related files.
+    """
+    save_path.mkdir(parents=True, exist_ok=True)
+
+    torch_model = TorchStaticModelPipeline(pipeline)
+    torch_model.eval()
+
+    # Persist the pipeline (encoder weights, head and config) alongside the ONNX graph
+    pipeline.save_pretrained(str(save_path))
+
+    # Prepare dummy input data
+    texts = ["hello", "hello world"]
+    input_ids, offsets = torch_model.tokenize(texts)
+
+    # An identity head is a regressor/projector, so its output is a raw value rather than a probability
+    output_name = "predictions" if pipeline.head.activation == Activation.IDENTITY else "probabilities"
+
+    onnx_model_path = save_path / "onnx/model.onnx"
+    onnx_model_path.parent.mkdir(parents=True, exist_ok=True)
+    torch.onnx.export(
+        torch_model,
+        (input_ids, offsets),
+        str(onnx_model_path),
+        export_params=True,
+        opset_version=14,
+        do_constant_folding=True,
+        input_names=["input_ids", "offsets"],
+        output_names=[output_name],
+        dynamic_axes={
+            "input_ids": {0: "num_tokens"},
+            "offsets": {0: "batch_size"},
+            output_name: {0: "batch_size"},
+        },
+    )
+
+    logger.info(f"Pipeline has been successfully exported to {onnx_model_path}")
+
+    # Save the tokenizer files required for transformers.js
+    save_tokenizer(pipeline.model.tokenizer, save_path)
+    logger.info(f"Tokenizer files have been saved to {save_path}")
+
+
 def save_tokenizer(tokenizer: Tokenizer, save_directory: Path) -> None:
     """Save tokenizer files in a format compatible with Transformers.
 
@@ -187,12 +293,12 @@ def save_tokenizer(tokenizer: Tokenizer, save_directory: Path) -> None:
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Export StaticModel to ONNX format")
+    parser = argparse.ArgumentParser(description="Export a StaticModel or StaticModelPipeline to ONNX format")
     parser.add_argument(
         "--model_path",
         type=str,
         required=True,
-        help="Path to the pretrained StaticModel",
+        help="Path to the pretrained StaticModel or StaticModelPipeline (classifier/regressor)",
     )
     parser.add_argument(
         "--save_path",

tests/test_export_to_onnx.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from model2vec.inference import StaticModelPipeline
+from model2vec.inference.mlp import Activation
+
+# The exporter lives in scripts/, which is not an installed package.
+sys.path.insert(0, str(Path(__file__).parent.parent / "scripts"))
+
+torch = pytest.importorskip("torch")
+ort = pytest.importorskip("onnxruntime")
+
+from export_to_onnx import TorchStaticModelPipeline  # noqa: E402
+
+
+def _export(torch_model: TorchStaticModelPipeline, texts: list[str], path: Path) -> np.ndarray:
+    """Export the wrapped pipeline to ONNX and run inference on the given texts."""
+    torch_model.eval()
+    input_ids, offsets = torch_model.tokenize(texts)
+    torch.onnx.export(
+        torch_model,
+        (input_ids, offsets),
+        str(path),
+        opset_version=14,
+        input_names=["input_ids", "offsets"],
+        output_names=["output"],
+        dynamic_axes={"input_ids": {0: "num_tokens"}, "offsets": {0: "batch_size"}, "output": {0: "batch_size"}},
+    )
+    session = ort.InferenceSession(str(path))
+    return session.run(None, {"input_ids": input_ids.numpy(), "offsets": offsets.numpy()})[0]
+
+
+def test_pipeline_onnx_matches_predict_proba(mock_inference_pipeline: StaticModelPipeline, tmp_path: Path) -> None:
+    """The exported classifier ONNX graph should reproduce the pipeline's probabilities."""
+    texts = ["dog", "cat", "dog cat"]
+    torch_model = TorchStaticModelPipeline(mock_inference_pipeline)
+
+    onnx_output = _export(torch_model, texts, tmp_path / "model.onnx")
+    expected = mock_inference_pipeline.predict_proba(texts, use_multiprocessing=False)
+
+    assert onnx_output.shape == expected.shape
+    np.testing.assert_allclose(onnx_output, expected, atol=1e-4)
+
+
+def test_pipeline_onnx_matches_projector(
+    mock_inference_pipeline_projector: StaticModelPipeline, tmp_path: Path
+) -> None:
+    """An identity-head (regressor/projector) exports raw predictions, not probabilities."""
+    assert mock_inference_pipeline_projector.head.activation == Activation.IDENTITY
+    texts = ["dog", "cat"]
+    torch_model = TorchStaticModelPipeline(mock_inference_pipeline_projector)
+
+    onnx_output = _export(torch_model, texts, tmp_path / "model.onnx")
+    expected = mock_inference_pipeline_projector.predict(texts, use_multiprocessing=False)
+
+    assert onnx_output.shape == expected.shape
+    np.testing.assert_allclose(onnx_output, expected, atol=1e-4)