From 33fe3641739f9f5b41077fc1b3793bdbda9c5584 Mon Sep 17 00:00:00 2001 From: stephantul Date: Thu, 20 Aug 2026 09:40:46 +0200 Subject: [PATCH 1/3] feat: make onnx export a function --- model2vec/onnx.py | 253 +++++++++++++++++++++ pyproject.toml | 4 +- scripts/export_to_onnx.py | 311 -------------------------- test.py | 15 ++ tests/integration/test_onnx_export.py | 80 +++++++ tests/test_export_to_onnx.py | 113 ++++++++-- uv.lock | 140 ++++++++++++ 7 files changed, 579 insertions(+), 337 deletions(-) create mode 100644 model2vec/onnx.py delete mode 100644 scripts/export_to_onnx.py create mode 100644 test.py create mode 100644 tests/integration/test_onnx_export.py diff --git a/model2vec/onnx.py b/model2vec/onnx.py new file mode 100644 index 0000000..cd6885c --- /dev/null +++ b/model2vec/onnx.py @@ -0,0 +1,253 @@ +from __future__ import annotations + +from model2vec.utils import get_package_extras, importable + +_REQUIRED_EXTRA = "onnx" + +for extra_dependency in get_package_extras("model2vec", _REQUIRED_EXTRA): + importable(extra_dependency, _REQUIRED_EXTRA) + +import json +import logging +import warnings +from pathlib import Path +from typing import Any + +import torch +from skeletoken import TokenizerModel +from tokenizers import Tokenizer +from torch.export import Dim + +from model2vec import StaticModel +from model2vec.inference import StaticModelPipeline +from model2vec.inference.mlp import Activation + +logger = logging.getLogger(__name__) + + +def _dummy_inputs() -> tuple[torch.Tensor, torch.Tensor]: + """Create dummy, padded (batch_size, sequence_length) `input_ids`/`attention_mask` tensors for tracing. + + :return: A tuple of (input_ids, attention_mask), both of shape (2, 3). + """ + input_ids = torch.zeros((2, 3), dtype=torch.long) + attention_mask = torch.tensor([[1, 1, 1], [1, 1, 0]], dtype=torch.long) + return input_ids, attention_mask + + +def _export_onnx(*args: Any, **kwargs: Any) -> None: + """Call `torch.onnx.export`, silencing the benign "axis name will not be used" warning. + + `input_ids` and `attention_mask` deliberately share the same `Dim` objects (see `_dynamic_shapes`), so + torch's exporter renames their shared symbol once and then warns, on the second input, that the (identical) + rename is redundant. That's expected here and not a sign of a problem. + + :param *args: Positional arguments forwarded to `torch.onnx.export`. + :param **kwargs: Keyword arguments forwarded to `torch.onnx.export`. + """ + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message=r".*will not be used, since it shares the same shape constraints") + torch.onnx.export(*args, **kwargs) + + +def _dynamic_shapes() -> dict[str, dict[int, Dim]]: + """Declare the dynamic (batch_size, sequence_length) axes shared by `input_ids` and `attention_mask`. + + :return: A `dynamic_shapes` mapping for `torch.onnx.export`. + """ + batch_size = Dim("batch_size") + sequence_length = Dim("sequence_length") + return { + "input_ids": {0: batch_size, 1: sequence_length}, + "attention_mask": {0: batch_size, 1: sequence_length}, + } + + +class TorchStaticModel(torch.nn.Module): + def __init__(self, model: StaticModel) -> None: + """Initialize the TorchStaticModel with a StaticModel instance.""" + super().__init__() + embeddings = torch.from_numpy(model.embedding) + if embeddings.dtype in {torch.int8, torch.uint8}: + embeddings = embeddings.to(torch.float16) + self.embeddings = torch.nn.Embedding.from_pretrained(embeddings, freeze=True) + self.normalize = model.normalize + + def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + """Forward pass of the model, following the standard transformers `input_ids`/`attention_mask` signature. + + :param input_ids: The input token ids, of shape (batch_size, sequence_length). + :param attention_mask: 1 for real tokens and 0 for padding, of shape (batch_size, sequence_length). + :return: The embeddings, of shape (batch_size, embedding_dim). + """ + mask = attention_mask.unsqueeze(-1).to(self.embeddings.weight.dtype) + # Zero out padding + embeddings = self.embeddings(input_ids) * mask + embeddings = embeddings.sum(dim=1) / mask.sum(dim=1).clamp(min=1) + # Normalize if required + if self.normalize: + embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=-1) + return embeddings + + +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, attention_mask: torch.Tensor) -> torch.Tensor: + """Encode the inputs and run them through the head, applying the output activation.""" + out = self.encoder(input_ids, attention_mask).float() + *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 export_model_to_onnx(model: StaticModel | StaticModelPipeline, save_path: str | Path) -> None: + """Export a StaticModel or a StaticModelPipeline to ONNX format and save tokenizer files. + + A classifier/regressor pipeline (one with a trained head) is exported with its head + included. A plain encoder is exported as embeddings. + + :param model: The StaticModel or StaticModelPipeline instance to export. + :param save_path: The directory to save the model and related files. + """ + save_path = Path(save_path) + if isinstance(model, StaticModelPipeline): + _export_pipeline_to_onnx(model, save_path) + return + + _export_encoder_to_onnx(model, save_path) + + +def _export_encoder_to_onnx(model: StaticModel, save_path: Path) -> None: + """Export a plain StaticModel encoder to ONNX format and save tokenizer files. + + :param model: The StaticModel instance to export. + :param save_path: The directory to save the model and related files. + """ + save_path.mkdir(parents=True, exist_ok=True) + + torch_model = TorchStaticModel(model) + torch_model.eval() + + # Prepare dummy input data, in the padded (batch_size, sequence_length) shape transformers-style ONNX runners expect + input_ids, attention_mask = _dummy_inputs() + + # Export the model to ONNX + onnx_model_path = save_path / "model.onnx" + onnx_model_path.parent.mkdir(parents=True, exist_ok=True) + _export_onnx( + torch_model, + (input_ids, attention_mask), + str(onnx_model_path), + export_params=True, + opset_version=18, + do_constant_folding=True, + input_names=["input_ids", "attention_mask"], + output_names=["embeddings"], + dynamic_shapes=_dynamic_shapes(), + ) + + logger.info(f"Model has been successfully exported to {onnx_model_path}") + + # Save the tokenizer files required for transformers.js, and a config.json for ONNX runtime providers + _save_tokenizer_and_config(model.tokenizer, save_path) + 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() + + # Prepare dummy input data, in the padded (batch_size, sequence_length) shape transformers-style ONNX runners expect + input_ids, attention_mask = _dummy_inputs() + + # 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 / "model.onnx" + onnx_model_path.parent.mkdir(parents=True, exist_ok=True) + _export_onnx( + torch_model, + (input_ids, attention_mask), + str(onnx_model_path), + export_params=True, + opset_version=18, + do_constant_folding=True, + input_names=["input_ids", "attention_mask"], + output_names=[output_name], + dynamic_shapes=_dynamic_shapes(), + ) + + logger.info(f"Pipeline has been successfully exported to {onnx_model_path}") + + # Save the tokenizer files required for transformers.js, and a config.json for ONNX runtime providers + _save_tokenizer_and_config(pipeline.model.tokenizer, save_path) + logger.info(f"Tokenizer files have been saved to {save_path}") + + +def _resolve_pad_token_id(tokenizer: Tokenizer, tokenizer_model: TokenizerModel) -> int: + """Resolve a pad token id for a StaticModel tokenizer, which may not have one registered as special. + + :param tokenizer: The tokenizer from the StaticModel. + :param tokenizer_model: `tokenizer` wrapped as a `skeletoken.TokenizerModel`. + :return: A vocabulary id to use as `pad_token_id`. + """ + if tokenizer_model.pad_token_id is not None: + return tokenizer_model.pad_token_id + + literal_pad_id = tokenizer.token_to_id("[PAD]") + if literal_pad_id is not None: + return literal_pad_id + + if tokenizer_model.unk_token_id is not None: + return tokenizer_model.unk_token_id + + return 0 + + +def _save_tokenizer_and_config(tokenizer: Tokenizer, save_directory: Path) -> None: + """Save tokenizer files in a format compatible with Transformers, plus config.json and special_tokens_map.json. + + :param tokenizer: The tokenizer from the StaticModel. + :param save_directory: The directory to save the tokenizer and config files. + """ + tokenizer_model = TokenizerModel.from_tokenizer(tokenizer) + tokenizer_model.to_transformers().save_pretrained(save_directory) + + pad_token_id = _resolve_pad_token_id(tokenizer, tokenizer_model) + pad_token = tokenizer.id_to_token(pad_token_id) or "" + + config = {"pad_token_id": pad_token_id} + (save_directory / "config.json").write_text(json.dumps(config, indent=2)) + + special_tokens_map = {"pad_token": pad_token} + (save_directory / "special_tokens_map.json").write_text(json.dumps(special_tokens_map, indent=2)) diff --git a/pyproject.toml b/pyproject.toml index aa24694..ec12d3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,8 +59,8 @@ dev = [ "setuptools", ] -distill = ["torch", "transformers<5.4.0", "skeletoken>=0.4.1,<0.5.0"] -onnx = ["onnx", "torch"] +distill = ["torch", "transformers<5.4.0", "skeletoken>=0.5.0"] +onnx = ["onnx", "torch", "onnxruntime", "onnxscript", "skeletoken>=0.5.0"] train = ["torch"] quantization = ["scikit-learn"] integration = ["mteb"] diff --git a/scripts/export_to_onnx.py b/scripts/export_to_onnx.py deleted file mode 100644 index b087045..0000000 --- a/scripts/export_to_onnx.py +++ /dev/null @@ -1,311 +0,0 @@ -from model2vec.utils import get_package_extras, importable - -# Define the optional dependency group name -_REQUIRED_EXTRA = "onnx" - -# Check if each dependency for the "onnx" group is importable -for extra_dependency in get_package_extras("model2vec", _REQUIRED_EXTRA): - importable(extra_dependency, _REQUIRED_EXTRA) - -import argparse -import json -import logging -from pathlib import Path - -import torch -from tokenizers import Tokenizer -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__) - - -class TorchStaticModel(torch.nn.Module): - def __init__(self, model: StaticModel) -> None: - """Initialize the TorchStaticModel with a StaticModel instance.""" - super().__init__() - # Convert NumPy embeddings to a torch.nn.EmbeddingBag - embeddings = torch.from_numpy(model.embedding) - if embeddings.dtype in {torch.int8, torch.uint8}: - embeddings = embeddings.to(torch.float16) - self.embedding_bag = torch.nn.EmbeddingBag.from_pretrained(embeddings, mode="mean", freeze=True) - self.normalize = model.normalize - # Save tokenizer attributes - self.tokenizer = model.tokenizer - self.unk_token_id = model.unk_token_id - self.median_token_length = model.median_token_length - - def forward(self, input_ids: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor: - """Forward pass of the model. - - :param input_ids: The input token ids. - :param offsets: The offsets to compute the mean pooling. - :return: The embeddings. - """ - # Perform embedding lookup and mean pooling - embeddings = self.embedding_bag(input_ids, offsets) - # Normalize if required - if self.normalize: - embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=-1) - return embeddings - - def tokenize(self, sentences: list[str], max_length: int | None = None) -> tuple[torch.Tensor, torch.Tensor]: - """Tokenize the input sentences. - - :param sentences: The input sentences. - :param max_length: The maximum length of the input_ids. - :return: The input_ids and offsets. - """ - # Tokenization logic similar to your StaticModel - if max_length is not None: - m = max_length * self.median_token_length - sentences = [sentence[:m] for sentence in sentences] - encodings = self.tokenizer.encode_batch(sentences, add_special_tokens=False) - encodings_ids = [encoding.ids for encoding in encodings] - if self.unk_token_id is not None: - # Remove unknown tokens - encodings_ids = [ - [token_id for token_id in token_ids if token_id != self.unk_token_id] for token_ids in encodings_ids - ] - if max_length is not None: - encodings_ids = [token_ids[:max_length] for token_ids in encodings_ids] - # Flatten input_ids and compute offsets - offsets = torch.tensor([0] + [len(ids) for ids in encodings_ids[:-1]], dtype=torch.long).cumsum(dim=0) - input_ids = torch.tensor( - [token_id for token_ids in encodings_ids for token_id in token_ids], - dtype=torch.long, - ) - 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 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. - """ - save_path.mkdir(parents=True, exist_ok=True) - - # Load the StaticModel - model = StaticModel.from_pretrained(model_path) - torch_model = TorchStaticModel(model) - - # Save the model using save_pretrained - model.save_pretrained(save_path) - - # Prepare dummy input data - texts = ["hello", "hello world"] - input_ids, offsets = torch_model.tokenize(texts) - - # Export the model to ONNX - 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=["embeddings"], - dynamic_axes={ - "input_ids": {0: "num_tokens"}, - "offsets": {0: "batch_size"}, - "embeddings": {0: "batch_size"}, - }, - ) - - logger.info(f"Model has been successfully exported to {onnx_model_path}") - - # Save the tokenizer files required for transformers.js - save_tokenizer(model.tokenizer, save_path) - 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. - - :param tokenizer: The tokenizer from the StaticModel. - :param save_directory: The directory to save the tokenizer files. - :raises FileNotFoundError: If config.json is not found in save_directory. - :raises ValueError: If tokenizer_name is not found in config.json. - """ - tokenizer_json_path = save_directory / "tokenizer.json" - tokenizer.save(str(tokenizer_json_path)) - - # Save vocab.txt - vocab = tokenizer.get_vocab() - vocab_path = save_directory / "vocab.txt" - with open(vocab_path, "w", encoding="utf-8") as vocab_file: - for token in sorted(vocab, key=vocab.get): - vocab_file.write(f"{token}\n") - - # Load config.json to get tokenizer_name - config_path = save_directory / "config.json" - if config_path.exists(): - with open(config_path, "r", encoding="utf-8") as f: - config = json.load(f) - else: - raise FileNotFoundError(f"config.json not found in {save_directory}") - - tokenizer_name = config.get("tokenizer_name") - if not tokenizer_name: - raise ValueError("tokenizer_name not found in config.json") - - # Load the original tokenizer - original_tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) - - # Extract special tokens and tokenizer class - special_tokens = original_tokenizer.special_tokens_map - tokenizer_class = original_tokenizer.__class__.__name__ - - # Load the tokenizer using PreTrainedTokenizerFast with special tokens - fast_tokenizer = PreTrainedTokenizerFast( - tokenizer_file=str(tokenizer_json_path), - **special_tokens, - ) - - # Save the tokenizer files - fast_tokenizer.save_pretrained(str(save_directory)) - # Modify tokenizer_config.json to set the correct tokenizer_class - tokenizer_config_path = save_directory / "tokenizer_config.json" - if tokenizer_config_path.exists(): - with open(tokenizer_config_path, "r", encoding="utf-8") as f: - tokenizer_config = json.load(f) - else: - raise FileNotFoundError(f"tokenizer_config.json not found in {save_directory}") - - # Update the tokenizer_class field - tokenizer_config["tokenizer_class"] = tokenizer_class - - # Write the updated tokenizer_config.json back to disk - with open(tokenizer_config_path, "w", encoding="utf-8") as f: - json.dump(tokenizer_config, f, indent=4, sort_keys=True) - - -if __name__ == "__main__": - 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 or StaticModelPipeline (classifier/regressor)", - ) - parser.add_argument( - "--save_path", - type=str, - required=True, - help="Directory to save the exported model and files", - ) - args = parser.parse_args() - - export_model_to_onnx(args.model_path, Path(args.save_path)) diff --git a/test.py b/test.py new file mode 100644 index 0000000..f643296 --- /dev/null +++ b/test.py @@ -0,0 +1,15 @@ +from fastembed import TextEmbedding +from fastembed.common.model_description import ModelSource, PoolingType + +model_name = "onnxmodel" + +TextEmbedding.add_custom_model( + model=model_name, + pooling=PoolingType.DISABLED, + normalization=True, + sources=ModelSource(url="onnxmodel"), + dim=256, + model_file="model.onnx", +) + +model = TextEmbedding(model_name=model_name, threads=1, specific_model_path="onnxmodel") diff --git a/tests/integration/test_onnx_export.py b/tests/integration/test_onnx_export.py new file mode 100644 index 0000000..e49fd7f --- /dev/null +++ b/tests/integration/test_onnx_export.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import onnxruntime as ort +import pytest +from skeletoken import TokenizerModel + +from model2vec import StaticModel +from model2vec.onnx import _resolve_pad_token_id, export_model_to_onnx + +_MODEL_NAME = "minishlab/potion-base-32m" + +_SENTENCES = [ + "The quick brown fox jumps over the lazy dog.", + "Paris is the capital of France.", + "I would like to order a large pizza with extra cheese.", + "Machine learning models can be trained on large datasets.", + "The weather today is sunny with a light breeze.", + "She sells seashells by the seashore.", + "This is a much longer sentence than the others, so padding and truncation are also exercised.", + "hi", +] + + +@pytest.fixture(scope="module") +def model() -> StaticModel: + """Download the real `potion-base-32m` model once for this module.""" + return StaticModel.from_pretrained(_MODEL_NAME) + + +@pytest.fixture(scope="module") +def export_dir(model: StaticModel, tmp_path_factory: pytest.TempPathFactory) -> Path: + """Export the real model to ONNX once, for reuse across every test in this module.""" + save_path = tmp_path_factory.mktemp("potion_onnx_export") + export_model_to_onnx(model, save_path) + return save_path + + +@pytest.fixture(scope="module") +def onnx_session(export_dir: Path) -> ort.InferenceSession: + """Load the exported ONNX graph into an onnxruntime session.""" + return ort.InferenceSession(str(export_dir / "model.onnx")) + + +def _tokenize(model: StaticModel, texts: list[str]) -> tuple[np.ndarray, np.ndarray]: + """Pad-tokenize texts into (input_ids, attention_mask) arrays, transformers-style.""" + tokenizer = model.tokenizer + pad_token_id = _resolve_pad_token_id(tokenizer, TokenizerModel.from_tokenizer(tokenizer)) + tokenizer.enable_padding(pad_id=pad_token_id) + encodings = tokenizer.encode_batch(texts, add_special_tokens=False) + tokenizer.no_padding() + input_ids = np.array([e.ids for e in encodings], dtype=np.int64) + attention_mask = np.array([e.attention_mask for e in encodings], dtype=np.int64) + return input_ids, attention_mask + + +def test_onnx_export_matches_encode_on_real_sentences(model: StaticModel, onnx_session: ort.InferenceSession) -> None: + """The ONNX graph exported from a real model should reproduce `StaticModel.encode` on real sentences.""" + input_ids, attention_mask = _tokenize(model, _SENTENCES) + + onnx_output = onnx_session.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})[0] + assert isinstance(onnx_output, np.ndarray) + expected = model.encode(_SENTENCES, use_multiprocessing=False) + + assert onnx_output.shape == expected.shape + np.testing.assert_allclose(onnx_output, expected, rtol=1e-4, atol=1e-4) + + +def test_onnx_export_writes_pad_token_id(model: StaticModel, export_dir: Path) -> None: + """The exported config.json and special_tokens_map.json must carry the real model's pad token.""" + expected_pad_token_id = _resolve_pad_token_id(model.tokenizer, TokenizerModel.from_tokenizer(model.tokenizer)) + + config = json.loads((export_dir / "config.json").read_text()) + assert config["pad_token_id"] == expected_pad_token_id + + special_tokens_map = json.loads((export_dir / "special_tokens_map.json").read_text()) + assert special_tokens_map["pad_token"] == model.tokenizer.id_to_token(expected_pad_token_id) diff --git a/tests/test_export_to_onnx.py b/tests/test_export_to_onnx.py index 9d0fe1b..374f6f8 100644 --- a/tests/test_export_to_onnx.py +++ b/tests/test_export_to_onnx.py @@ -1,52 +1,116 @@ from __future__ import annotations -import sys +import json from pathlib import Path import numpy as np -import pytest - +import onnxruntime as ort +import torch +from skeletoken import TokenizerModel +from tokenizers import Tokenizer +from tokenizers.models import BPE +from tokenizers.pre_tokenizers import Whitespace + +from model2vec import StaticModel 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.""" +from model2vec.onnx import ( + TorchStaticModelPipeline, + _dynamic_shapes, + _export_onnx, + _resolve_pad_token_id, + export_model_to_onnx, +) + + +def _tokenize(pipeline: StaticModelPipeline, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]: + """Pad-tokenize texts into (input_ids, attention_mask) tensors, transformers-style.""" + tokenizer = pipeline.model.tokenizer + tokenizer.enable_padding(pad_id=0, pad_token="[PAD]") + encodings = tokenizer.encode_batch(texts, add_special_tokens=False) + tokenizer.no_padding() + input_ids = torch.tensor([e.ids for e in encodings], dtype=torch.long) + attention_mask = torch.tensor([e.attention_mask for e in encodings], dtype=torch.long) + return input_ids, attention_mask + + +def _export( + torch_model: TorchStaticModelPipeline, input_ids: torch.Tensor, attention_mask: torch.Tensor, path: Path +) -> np.ndarray: + """Export the wrapped pipeline to ONNX and run inference on the given inputs.""" torch_model.eval() - input_ids, offsets = torch_model.tokenize(texts) - torch.onnx.export( + _export_onnx( torch_model, - (input_ids, offsets), + (input_ids, attention_mask), str(path), - opset_version=14, - input_names=["input_ids", "offsets"], + opset_version=18, + input_names=["input_ids", "attention_mask"], output_names=["output"], - dynamic_axes={"input_ids": {0: "num_tokens"}, "offsets": {0: "batch_size"}, "output": {0: "batch_size"}}, + dynamic_shapes=_dynamic_shapes(), ) session = ort.InferenceSession(str(path)) - return session.run(None, {"input_ids": input_ids.numpy(), "offsets": offsets.numpy()})[0] + output = session.run(None, {"input_ids": input_ids.numpy(), "attention_mask": attention_mask.numpy()})[0] + assert isinstance(output, np.ndarray) + return output 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) + input_ids, attention_mask = _tokenize(mock_inference_pipeline, texts) - onnx_output = _export(torch_model, texts, tmp_path / "model.onnx") + onnx_output = _export(torch_model, input_ids, attention_mask, 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_export_model_to_onnx_encoder(mock_static_model: StaticModel, tmp_path: Path) -> None: + """A plain StaticModel is exported with a config.json and special_tokens_map.json alongside it.""" + save_path = tmp_path / "export" + export_model_to_onnx(mock_static_model, save_path) + + assert (save_path / "model.onnx").exists() + tokenizer_model = TokenizerModel.from_tokenizer(mock_static_model.tokenizer) + expected_pad_token_id = _resolve_pad_token_id(mock_static_model.tokenizer, tokenizer_model) + + config = json.loads((save_path / "config.json").read_text()) + assert config["pad_token_id"] == expected_pad_token_id + + special_tokens_map = json.loads((save_path / "special_tokens_map.json").read_text()) + assert special_tokens_map["pad_token"] == mock_static_model.tokenizer.id_to_token(expected_pad_token_id) + + +def test_export_model_to_onnx_pipeline(mock_inference_pipeline: StaticModelPipeline, tmp_path: Path) -> None: + """A StaticModelPipeline is exported with a config.json and special_tokens_map.json alongside it.""" + save_path = tmp_path / "export" + export_model_to_onnx(mock_inference_pipeline, save_path) + + assert (save_path / "model.onnx").exists() + tokenizer_model = TokenizerModel.from_tokenizer(mock_inference_pipeline.model.tokenizer) + expected_pad_token_id = _resolve_pad_token_id(mock_inference_pipeline.model.tokenizer, tokenizer_model) + + config = json.loads((save_path / "config.json").read_text()) + assert config["pad_token_id"] == expected_pad_token_id + + special_tokens_map = json.loads((save_path / "special_tokens_map.json").read_text()) + assert special_tokens_map["pad_token"] == mock_inference_pipeline.model.tokenizer.id_to_token(expected_pad_token_id) + + +def test_resolve_pad_token_id_falls_back_to_unk_when_no_pad_registered() -> None: + """When a tokenizer has no registered pad token and no literal "[PAD]" entry, fall back to unk, not position 0.""" + vocab = ["!", "hello", "world", "[UNK]"] + tokenizer = Tokenizer(BPE(vocab={t: i for i, t in enumerate(vocab)}, merges=[], unk_token="[UNK]")) + tokenizer.pre_tokenizer = Whitespace() # type: ignore[assignment] + tokenizer_model = TokenizerModel.from_tokenizer(tokenizer) + assert tokenizer_model.pad_token_id is None + assert tokenizer.token_to_id("[PAD]") is None + + assert _resolve_pad_token_id(tokenizer, tokenizer_model) == tokenizer.token_to_id("[UNK]") + + def test_pipeline_onnx_matches_projector( mock_inference_pipeline_projector: StaticModelPipeline, tmp_path: Path ) -> None: @@ -54,8 +118,9 @@ def test_pipeline_onnx_matches_projector( assert mock_inference_pipeline_projector.head.activation == Activation.IDENTITY texts = ["dog", "cat"] torch_model = TorchStaticModelPipeline(mock_inference_pipeline_projector) + input_ids, attention_mask = _tokenize(mock_inference_pipeline_projector, texts) - onnx_output = _export(torch_model, texts, tmp_path / "model.onnx") + onnx_output = _export(torch_model, input_ids, attention_mask, tmp_path / "model.onnx") expected = mock_inference_pipeline_projector.predict(texts, use_multiprocessing=False) assert onnx_output.shape == expected.shape diff --git a/uv.lock b/uv.lock index 8577209..52d3018 100644 --- a/uv.lock +++ b/uv.lock @@ -595,6 +595,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a4/a5/842ae8f0c08b61d6484b52f99a03510a3a72d23141942d216ebe81fefbce/filelock-3.25.2-py3-none-any.whl", hash = "sha256:ca8afb0da15f229774c9ad1b455ed96e85a81373065fb10446672f64444ddf70", size = 26759, upload-time = "2026-03-11T20:45:37.437Z" }, ] +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl", hash = "sha256:7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4", size = 26661, upload-time = "2025-12-19T23:16:13.622Z" }, +] + [[package]] name = "frozenlist" version = "1.8.0" @@ -1260,6 +1268,9 @@ integration = [ ] onnx = [ { name = "onnx" }, + { name = "onnxruntime", version = "1.24.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "onnxruntime", version = "1.28.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "onnxscript" }, { name = "torch" }, ] quantization = [ @@ -1285,6 +1296,8 @@ requires-dist = [ { name = "mypy", marker = "extra == 'dev'" }, { name = "numpy" }, { name = "onnx", marker = "extra == 'onnx'" }, + { name = "onnxruntime", marker = "extra == 'onnx'" }, + { name = "onnxscript", marker = "extra == 'onnx'" }, { name = "pre-commit", marker = "extra == 'dev'" }, { name = "pytest", marker = "extra == 'dev'" }, { name = "pytest-cov", marker = "extra == 'dev'" }, @@ -1930,6 +1943,133 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/aa/7d/1bbe626ff6b192c844d3ad34356840cc60fca02e2dea0db95e01645758b1/onnx-1.20.1-cp313-cp313t-win_arm64.whl", hash = "sha256:eb335d7bcf9abac82a0d6a0fda0363531ae0b22cfd0fc6304bff32ee29905def", size = 16348968, upload-time = "2026-01-10T01:40:00.491Z" }, ] +[[package]] +name = "onnx-ir" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "onnx" }, + { name = "sympy" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d6/c2/61194cec0dbc5622273c0ebd592d37cc1dca0d7f1a744f02edd45ac905a3/onnx_ir-1.0.0.tar.gz", hash = "sha256:9e261f25fde8da9612ae5cb43b3b374d5ff469c04af0363cad588b2bb000b812", size = 163121, upload-time = "2026-08-11T14:49:46.895Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/cd/6d1637172eb59c7b18ac90ed089d1f599a11fe0e63b4db2d017f3bb38a32/onnx_ir-1.0.0-py3-none-any.whl", hash = "sha256:e578f0d608d3062866b48223616eb2d10a6d6d01f8b8faac596129034f483cc7", size = 185849, upload-time = "2026-08-11T14:49:45.524Z" }, +] + +[[package]] +name = "onnxruntime" +version = "1.24.3" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] +dependencies = [ + { name = "flatbuffers" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "sympy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/15/41/3253db975a90c3ce1d475e2a230773a21cd7998537f0657947df6fb79861/onnxruntime-1.24.3-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:3e6456801c66b095c5cd68e690ca25db970ea5202bd0c5b84a2c3ef7731c5a3c", size = 17332766, upload-time = "2026-03-05T17:18:59.714Z" }, + { url = "https://files.pythonhosted.org/packages/7e/c5/3af6b325f1492d691b23844d88ed26844c1164620860c5efe95c0e22782d/onnxruntime-1.24.3-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8b2ebc54c6d8281dccff78d4b06e47d4cf07535937584ab759448390a70f4978", size = 15130330, upload-time = "2026-03-05T16:34:53.831Z" }, + { url = "https://files.pythonhosted.org/packages/03/4b/f96b46c1866a293ed23ca2cf5e5a63d413ad3a951da60dd877e3c56cbbca/onnxruntime-1.24.3-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fb56575d7794bf0781156955610c9e651c9504c64d42ec880784b6106244882d", size = 17213247, upload-time = "2026-03-05T17:17:59.812Z" }, + { url = "https://files.pythonhosted.org/packages/36/13/27cf4d8df2578747584e8758aeb0b673b60274048510257f1f084b15e80e/onnxruntime-1.24.3-cp311-cp311-win_amd64.whl", hash = "sha256:c958222ef9eff54018332beecd32d5d94a3ab079d8821937b333811bf4da0d39", size = 12595530, upload-time = "2026-03-05T17:18:49.356Z" }, + { url = "https://files.pythonhosted.org/packages/19/8c/6d9f31e6bae72a8079be12ed8ba36c4126a571fad38ded0a1b96f60f6896/onnxruntime-1.24.3-cp311-cp311-win_arm64.whl", hash = "sha256:a8f761857ebaf58a85b9e42422d03207f1d39e6bb8fecfdbf613bac5b9710723", size = 12261715, upload-time = "2026-03-05T17:18:39.699Z" }, + { url = "https://files.pythonhosted.org/packages/d0/7f/dfdc4e52600fde4c02d59bfe98c4b057931c1114b701e175aee311a9bc11/onnxruntime-1.24.3-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:0d244227dc5e00a9ae15a7ac1eba4c4460d7876dfecafe73fb00db9f1d914d91", size = 17342578, upload-time = "2026-03-05T17:19:02.403Z" }, + { url = "https://files.pythonhosted.org/packages/1c/dc/1f5489f7b21817d4ad352bf7a92a252bd5b438bcbaa7ad20ea50814edc79/onnxruntime-1.24.3-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0a9847b870b6cb462652b547bc98c49e0efb67553410a082fde1918a38707452", size = 15150105, upload-time = "2026-03-05T16:34:56.897Z" }, + { url = "https://files.pythonhosted.org/packages/28/7c/fd253da53594ab8efbefdc85b3638620ab1a6aab6eb7028a513c853559ce/onnxruntime-1.24.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b354afce3333f2859c7e8706d84b6c552beac39233bcd3141ce7ab77b4cabb5d", size = 17237101, upload-time = "2026-03-05T17:18:02.561Z" }, + { url = "https://files.pythonhosted.org/packages/71/5f/eaabc5699eeed6a9188c5c055ac1948ae50138697a0428d562ac970d7db5/onnxruntime-1.24.3-cp312-cp312-win_amd64.whl", hash = "sha256:44ea708c34965439170d811267c51281d3897ecfc4aa0087fa25d4a4c3eb2e4a", size = 12597638, upload-time = "2026-03-05T17:18:52.141Z" }, + { url = "https://files.pythonhosted.org/packages/cc/5c/d8066c320b90610dbeb489a483b132c3b3879b2f93f949fb5d30cfa9b119/onnxruntime-1.24.3-cp312-cp312-win_arm64.whl", hash = "sha256:48d1092b44ca2ba6f9543892e7c422c15a568481403c10440945685faf27a8d8", size = 12270943, upload-time = "2026-03-05T17:18:42.006Z" }, + { url = "https://files.pythonhosted.org/packages/51/8d/487ece554119e2991242d4de55de7019ac6e47ee8dfafa69fcf41d37f8ed/onnxruntime-1.24.3-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:34a0ea5ff191d8420d9c1332355644148b1bf1a0d10c411af890a63a9f662aa7", size = 17342706, upload-time = "2026-03-05T16:35:10.813Z" }, + { url = "https://files.pythonhosted.org/packages/dd/25/8b444f463c1ac6106b889f6235c84f01eec001eaf689c3eff8c69cf48fae/onnxruntime-1.24.3-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1fd2ec7bb0fabe42f55e8337cfc9b1969d0d14622711aac73d69b4bd5abb5ed7", size = 15149956, upload-time = "2026-03-05T16:34:59.264Z" }, + { url = "https://files.pythonhosted.org/packages/34/fc/c9182a3e1ab46940dd4f30e61071f59eee8804c1f641f37ce6e173633fb6/onnxruntime-1.24.3-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:df8e70e732fe26346faaeec9147fa38bef35d232d2495d27e93dd221a2d473a9", size = 17237370, upload-time = "2026-03-05T17:18:05.258Z" }, + { url = "https://files.pythonhosted.org/packages/05/7e/3b549e1f4538514118bff98a1bcd6481dd9a17067f8c9af77151621c9a5c/onnxruntime-1.24.3-cp313-cp313-win_amd64.whl", hash = "sha256:2d3706719be6ad41d38a2250998b1d87758a20f6ea4546962e21dc79f1f1fd2b", size = 12597939, upload-time = "2026-03-05T17:18:54.772Z" }, + { url = "https://files.pythonhosted.org/packages/80/41/9696a5c4631a0caa75cc8bc4efd30938fd483694aa614898d087c3ee6d29/onnxruntime-1.24.3-cp313-cp313-win_arm64.whl", hash = "sha256:b082f3ba9519f0a1a1e754556bc7e635c7526ef81b98b3f78da4455d25f0437b", size = 12270705, upload-time = "2026-03-05T17:18:44.774Z" }, + { url = "https://files.pythonhosted.org/packages/b7/65/a26c5e59e3b210852ee04248cf8843c81fe7d40d94cf95343b66efe7eec9/onnxruntime-1.24.3-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:72f956634bc2e4bd2e8b006bef111849bd42c42dea37bd0a4c728404fdaf4d34", size = 15161796, upload-time = "2026-03-05T16:35:02.871Z" }, + { url = "https://files.pythonhosted.org/packages/f3/25/2035b4aa2ccb5be6acf139397731ec507c5f09e199ab39d3262b22ffa1ac/onnxruntime-1.24.3-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:78d1f25eed4ab9959db70a626ed50ee24cf497e60774f59f1207ac8556399c4d", size = 17240936, upload-time = "2026-03-05T17:18:09.534Z" }, + { url = "https://files.pythonhosted.org/packages/f9/a4/b3240ea84b92a3efb83d49cc16c04a17ade1ab47a6a95c4866d15bf0ac35/onnxruntime-1.24.3-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:a6b4bce87d96f78f0a9bf5cefab3303ae95d558c5bfea53d0bf7f9ea207880a8", size = 17344149, upload-time = "2026-03-05T16:35:13.382Z" }, + { url = "https://files.pythonhosted.org/packages/bb/4a/4b56757e51a56265e8c56764d9c36d7b435045e05e3b8a38bedfc5aedba3/onnxruntime-1.24.3-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d48f36c87b25ab3b2b4c88826c96cf1399a5631e3c2c03cc27d6a1e5d6b18eb4", size = 15151571, upload-time = "2026-03-05T16:35:05.679Z" }, + { url = "https://files.pythonhosted.org/packages/cf/14/c6fb84980cec8f682a523fcac7c2bdd6b311e7f342c61ce48d3a9cb87fc6/onnxruntime-1.24.3-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e104d33a409bf6e3f30f0e8198ec2aaf8d445b8395490a80f6e6ad56da98e400", size = 17238951, upload-time = "2026-03-05T17:18:12.394Z" }, + { url = "https://files.pythonhosted.org/packages/57/14/447e1400165aca8caf35dabd46540eb943c92f3065927bb4d9bcbc91e221/onnxruntime-1.24.3-cp314-cp314-win_amd64.whl", hash = "sha256:e785d73fbd17421c2513b0bb09eb25d88fa22c8c10c3f5d6060589efa5537c5b", size = 12903820, upload-time = "2026-03-05T17:18:57.123Z" }, + { url = "https://files.pythonhosted.org/packages/1d/ec/6b2fa5702e4bbba7339ca5787a9d056fc564a16079f8833cc6ba4798da1c/onnxruntime-1.24.3-cp314-cp314-win_arm64.whl", hash = "sha256:951e897a275f897a05ffbcaa615d98777882decaeb80c9216c68cdc62f849f53", size = 12594089, upload-time = "2026-03-05T17:18:47.169Z" }, + { url = "https://files.pythonhosted.org/packages/12/dc/cd06cba3ddad92ceb17b914a8e8d49836c79e38936e26bde6e368b62c1fe/onnxruntime-1.24.3-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4d4e70ce578aa214c74c7a7a9226bc8e229814db4a5b2d097333b81279ecde36", size = 15162789, upload-time = "2026-03-05T16:35:08.282Z" }, + { url = "https://files.pythonhosted.org/packages/a6/d6/413e98ab666c6fb9e8be7d1c6eb3bd403b0bea1b8d42db066dab98c7df07/onnxruntime-1.24.3-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:02aaf6ddfa784523b6873b4176a79d508e599efe12ab0ea1a3a6e7314408b7aa", size = 17240738, upload-time = "2026-03-05T17:18:15.203Z" }, +] + +[[package]] +name = "onnxruntime" +version = "1.28.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", +] +dependencies = [ + { name = "flatbuffers" }, + { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" } }, + { name = "packaging" }, + { name = "protobuf" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/17/4d/5014667e2a3a77d6e1b74cc3d88948d06163b8e0a33a84c85073322b5dec/onnxruntime-1.28.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:f5c5daabd28aad610f83fdcf32acec8fb57e6adc6c6a39fe2a3c755db957b410", size = 19130506, upload-time = "2026-07-25T01:22:34.489Z" }, + { url = "https://files.pythonhosted.org/packages/ea/97/b7ce1bc8bb6048b5fe9129f55d6506dc19499068ef2e0a0af1ae3c8aa4e7/onnxruntime-1.28.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8d66f9ceb29909c70839e4e4fb3435c7b490050d8f162bd5f3aba4ca01ee517f", size = 17039880, upload-time = "2026-07-25T01:21:37.538Z" }, + { url = "https://files.pythonhosted.org/packages/f3/17/4e5ecd8764f87573c495d834ce79e61ecca47f7a01d1e444a606e570edcb/onnxruntime-1.28.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a166b78ee04f3a37fa1ef82034b6a3ce96d9684e582d4d30b296de83e9998bb5", size = 19193162, upload-time = "2026-07-25T01:21:59.151Z" }, + { url = "https://files.pythonhosted.org/packages/9f/10/3d946d5d5f2cdcc3c8da36cae63190c516d16349edaffd944bda60ca4c3e/onnxruntime-1.28.0-cp311-cp311-win_amd64.whl", hash = "sha256:0d650aeee29368414367b65529e90afe4bf1bab76254789063b8b2f7ea3013c8", size = 13752539, upload-time = "2026-07-25T01:22:24.524Z" }, + { url = "https://files.pythonhosted.org/packages/8f/74/1c440be7af1e026280b139caa1be5d11bd4dc368011ddbe8f5362b58e12f/onnxruntime-1.28.0-cp311-cp311-win_arm64.whl", hash = "sha256:0faf85fb447a663c9cdadc39bd6b19bdf7bedded6699e45731b9b36c46fd993d", size = 13449940, upload-time = "2026-07-25T01:22:14.97Z" }, + { url = "https://files.pythonhosted.org/packages/98/f8/dcbe7700dca82fa540035abd3c868fe5ad0f86af00b9a3db7c2e27d15c7d/onnxruntime-1.28.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:26ff0fdd06efb6c155bae95387a09db1a2be89c7a03e4d0bffd5a171cc2826da", size = 19141362, upload-time = "2026-07-25T01:22:36.965Z" }, + { url = "https://files.pythonhosted.org/packages/28/5b/1d77e62097fdbe07e2dc827f389b1c4c0c275f6fab0369a8f46d2461af27/onnxruntime-1.28.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e81a23df16e7acb9d51b06d30cc098e49315ef9180f97bc2221d167b4b04d9c", size = 17050628, upload-time = "2026-07-25T01:21:40.481Z" }, + { url = "https://files.pythonhosted.org/packages/95/df/5486ab03e9be288d5268867054c8b04bebcf95bfd12e801c05cc67703dab/onnxruntime-1.28.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0a83bdb70d143cede762b677789bf2a7acca54b3fb82565601d5c30695aa933c", size = 19214257, upload-time = "2026-07-25T01:22:01.695Z" }, + { url = "https://files.pythonhosted.org/packages/3e/3b/986ca67c274932ba9ac5332fb10de56f643dfd433c74e33f8ae8f847cf24/onnxruntime-1.28.0-cp312-cp312-win_amd64.whl", hash = "sha256:c35064f9b3c43c81c5d5d282091401d0f1ff22796d93ccade4ea2ece5e137ab8", size = 13755036, upload-time = "2026-07-25T01:22:26.89Z" }, + { url = "https://files.pythonhosted.org/packages/1d/46/059dba81d46c6ba88e0c2d1c64321ac8098847d678423300a183d42ecbd6/onnxruntime-1.28.0-cp312-cp312-win_arm64.whl", hash = "sha256:e02feeb0165c5f13b4cc954738078d59b90128516ac12b671ee24a530242bf02", size = 13454462, upload-time = "2026-07-25T01:22:17.38Z" }, + { url = "https://files.pythonhosted.org/packages/9c/12/3807e2b17d9eb71d3cb78ed2ba76869b05c637c9b9d6112e636098b0c97a/onnxruntime-1.28.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:31410f544674f534c2f27348af52ef81682ca9c8719154bf4d48f0ef23823b1e", size = 19141759, upload-time = "2026-07-25T01:21:53.765Z" }, + { url = "https://files.pythonhosted.org/packages/c0/23/b46045c3bf67a9cf54c12f5df0f018a422c65fbb9d6072b10071bebfaae2/onnxruntime-1.28.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f649dd6f6452d12a8059888aa489fe519e062e18793dac72b9efa0f9fdb64135", size = 17049339, upload-time = "2026-07-25T01:21:43.005Z" }, + { url = "https://files.pythonhosted.org/packages/78/b6/8c5396e7894e77c5a7d1e026f3acb9dd39c4b5644e412e37a0055eaa3bc5/onnxruntime-1.28.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:54fa221d669282bd8f582708ce4c96010a7e9fb0661f9006b37fe2fedafb73fe", size = 19214329, upload-time = "2026-07-25T01:22:04.133Z" }, + { url = "https://files.pythonhosted.org/packages/56/f1/51225c202edba4dfc94e1ea03f3d78f1aaf307da75fd792c0ce1946b2514/onnxruntime-1.28.0-cp313-cp313-win_amd64.whl", hash = "sha256:1a1a19175464665c9b8d50bc916f216cc0b569110045b7bbca8f9f290b186f58", size = 13755033, upload-time = "2026-07-25T01:22:29.302Z" }, + { url = "https://files.pythonhosted.org/packages/f4/db/f59f715edfdd96a051f32b5ef0e680a20a8755d4ecd75f63090e960e347a/onnxruntime-1.28.0-cp313-cp313-win_arm64.whl", hash = "sha256:cfab507abe09d6ffeb817eee07944d452fdc0b00fdcef34cab4db10a45e378c7", size = 13454175, upload-time = "2026-07-25T01:22:19.912Z" }, + { url = "https://files.pythonhosted.org/packages/47/28/810314fa88647af9f4cdaf438a30ad1cfebebb53ded55499232d7a0094e6/onnxruntime-1.28.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ac301f53b1930402fc46c368e268acfed02f3207272aaff05070d7e09f96f031", size = 17057307, upload-time = "2026-07-25T01:21:45.492Z" }, + { url = "https://files.pythonhosted.org/packages/3d/cc/9e9f193cc0f29f263a8f09ec08487aed6c96ee856d5fd77da32a425c1949/onnxruntime-1.28.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f7f022a1103cae591c75fc4565589a515f2ddd14a6ac8e8a05812dfeda142e28", size = 19222954, upload-time = "2026-07-25T01:22:06.952Z" }, + { url = "https://files.pythonhosted.org/packages/4e/eb/952314c451d9463e5c9aed9978eec76cf32930d407d9ab8700dd0f4ea1ea/onnxruntime-1.28.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:8adff67a3f28257b37cfe945a7e952e4122666aa8c91a0380862e9fd4c2ed19f", size = 19143748, upload-time = "2026-07-25T01:21:56.297Z" }, + { url = "https://files.pythonhosted.org/packages/3f/e9/139180b4dd810329aaa42c238b4e6383c906202d98609ae29d66eb7c32b1/onnxruntime-1.28.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bc2565e487b4896fb988d6383577d875d958e071fc5f6c3550bd5d02ae98264b", size = 17051950, upload-time = "2026-07-25T01:21:48.606Z" }, + { url = "https://files.pythonhosted.org/packages/03/88/9432428273356ad3c8aa01f52c1b3e7f53c4c0192748f41ad983872b436b/onnxruntime-1.28.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6afdc83f1317c136e92fc29f5ee9f058de59d87c0b22cee3fdbfbaa0ccc2098a", size = 19214924, upload-time = "2026-07-25T01:22:09.727Z" }, + { url = "https://files.pythonhosted.org/packages/bb/e2/6feb3a43517aaf2b1bf7e46897ba5eb81a29717f7d7901420614d5ee4653/onnxruntime-1.28.0-cp314-cp314-win_amd64.whl", hash = "sha256:f2a3b9e30ce880d4ca54999cb313569e36da4f62eefe25f87be18f43e9a3a4d5", size = 14093738, upload-time = "2026-07-25T01:22:31.629Z" }, + { url = "https://files.pythonhosted.org/packages/fc/8f/83974a1e201dc2e58e5e7111bcaeb1ca2413e9c41f505d26419ee9e3dddf/onnxruntime-1.28.0-cp314-cp314-win_arm64.whl", hash = "sha256:07fb3cbe990d6bf0ab3c22bfbbfb0e314151266046ea6edb4a07f556b4258c5f", size = 13821117, upload-time = "2026-07-25T01:22:22.387Z" }, + { url = "https://files.pythonhosted.org/packages/0d/83/00e606bc25c756d76a267370c39b7516ad52f9cf134d7ff2bff8b6108bc4/onnxruntime-1.28.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e562d6e36a749f6764481c0ddb0f2af3d0b5a3c164291361d08803c557f369af", size = 17055518, upload-time = "2026-07-25T01:21:51.08Z" }, + { url = "https://files.pythonhosted.org/packages/94/a9/68707e1ce345cbdbcd4df65932ebc82a673e917d63eda0007ebcff948691/onnxruntime-1.28.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4f6e92367ddce1e4d33cf295024f40192be6c6171a09208f515ba169ced06c8e", size = 19222976, upload-time = "2026-07-25T01:22:12.474Z" }, +] + +[[package]] +name = "onnxscript" +version = "0.7.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "onnx" }, + { name = "onnx-ir" }, + { name = "packaging" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c4/3a/4d79bce3f460e0df7fed54a92ce80827f25da66511da368bb00783ad8d20/onnxscript-0.7.1.tar.gz", hash = "sha256:309fb86484b11fa4ded90dba580e0d63f1a0827588e521cecaf2eeddb46d6e86", size = 618160, upload-time = "2026-06-29T23:33:21.526Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dd/bd/a0c8e737b6afda10e42a597787d53d5b66e00268df6f59184701eeae37d9/onnxscript-0.7.1-py3-none-any.whl", hash = "sha256:544763b7fdef49940cdd9412ff5135cbae96d59ac6bc1921457f21280f40f4b7", size = 721970, upload-time = "2026-06-29T23:33:23.298Z" }, +] + [[package]] name = "packaging" version = "26.0" From f625b0e89306bf66076a1d2d556db64918f41b34 Mon Sep 17 00:00:00 2001 From: stephantul Date: Thu, 20 Aug 2026 14:37:17 +0200 Subject: [PATCH 2/3] remove test file --- test.py | 15 --------------- 1 file changed, 15 deletions(-) delete mode 100644 test.py diff --git a/test.py b/test.py deleted file mode 100644 index f643296..0000000 --- a/test.py +++ /dev/null @@ -1,15 +0,0 @@ -from fastembed import TextEmbedding -from fastembed.common.model_description import ModelSource, PoolingType - -model_name = "onnxmodel" - -TextEmbedding.add_custom_model( - model=model_name, - pooling=PoolingType.DISABLED, - normalization=True, - sources=ModelSource(url="onnxmodel"), - dim=256, - model_file="model.onnx", -) - -model = TextEmbedding(model_name=model_name, threads=1, specific_model_path="onnxmodel") From 6e464173051ba268e80947478d3e933d5569c461 Mon Sep 17 00:00:00 2001 From: stephantul Date: Thu, 20 Aug 2026 15:49:49 +0200 Subject: [PATCH 3/3] set external data to false --- model2vec/onnx.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/model2vec/onnx.py b/model2vec/onnx.py index cd6885c..55236f5 100644 --- a/model2vec/onnx.py +++ b/model2vec/onnx.py @@ -164,6 +164,7 @@ def _export_encoder_to_onnx(model: StaticModel, save_path: Path) -> None: input_names=["input_ids", "attention_mask"], output_names=["embeddings"], dynamic_shapes=_dynamic_shapes(), + external_data=False, ) logger.info(f"Model has been successfully exported to {onnx_model_path}") @@ -205,6 +206,7 @@ def _export_pipeline_to_onnx(pipeline: StaticModelPipeline, save_path: Path) -> input_names=["input_ids", "attention_mask"], output_names=[output_name], dynamic_shapes=_dynamic_shapes(), + external_data=False, ) logger.info(f"Pipeline has been successfully exported to {onnx_model_path}") @@ -240,11 +242,18 @@ def _save_tokenizer_and_config(tokenizer: Tokenizer, save_directory: Path) -> No :param tokenizer: The tokenizer from the StaticModel. :param save_directory: The directory to save the tokenizer and config files. """ + """with TemporaryDirectory() as tmp: + tokenizer.save(str(Path(tmp) / "tokenizer.json")) + tokenizer_model = TokenizerModel.from_pretrained(Path(tmp) / "tokenizer.json")""" tokenizer_model = TokenizerModel.from_tokenizer(tokenizer) - tokenizer_model.to_transformers().save_pretrained(save_directory) - pad_token_id = _resolve_pad_token_id(tokenizer, tokenizer_model) pad_token = tokenizer.id_to_token(pad_token_id) or "" + if pad_token: + tokenizer_model.pad_token = pad_token + hf = tokenizer_model.to_transformers() + # Hardcoded max length of 32768 + hf.model_max_length = 32768 + hf.save_pretrained(save_directory) config = {"pad_token_id": pad_token_id} (save_directory / "config.json").write_text(json.dumps(config, indent=2))