diff --git a/scripts/export_to_onnx.py b/scripts/export_to_onnx.py index dda4c5a6..b0870455 100644 --- a/scripts/export_to_onnx.py +++ b/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", diff --git a/tests/test_export_to_onnx.py b/tests/test_export_to_onnx.py new file mode 100644 index 00000000..9d0fe1b5 --- /dev/null +++ b/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)