From 93acf3d669bcbbd7c0c644c0e1921775ebd274cb Mon Sep 17 00:00:00 2001 From: serhiizghama Date: Fri, 11 Sep 2026 20:04:40 +0700 Subject: [PATCH 1/2] fix(onnx): export vocabulary-quantized models with their token mapping and weights --- model2vec/onnx.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/model2vec/onnx.py b/model2vec/onnx.py index 37a719b..73900e7 100644 --- a/model2vec/onnx.py +++ b/model2vec/onnx.py @@ -65,6 +65,9 @@ def _dynamic_shapes() -> dict[str, dict[int, Dim]]: class TorchStaticModel(torch.nn.Module): + token_mapping: torch.Tensor | None + weights: torch.Tensor | None + def __init__(self, model: StaticModel) -> None: """Initialize the TorchStaticModel with a StaticModel instance.""" super().__init__() @@ -72,6 +75,11 @@ def __init__(self, model: StaticModel) -> None: if embeddings.dtype in {torch.int8, torch.uint8}: embeddings = embeddings.to(torch.float16) self.embeddings = torch.nn.Embedding.from_pretrained(embeddings, freeze=True) + # Vocabulary-quantized models look up ids through `token_mapping` and weigh tokens, like `_encode_helper` + token_mapping = None if model.token_mapping is None else torch.from_numpy(model.token_mapping).long() + weights = None if model.weights is None else torch.from_numpy(model.weights).to(embeddings.dtype) + self.register_buffer("token_mapping", token_mapping) + self.register_buffer("weights", weights) self.normalize = model.normalize self.unk_token_id = model.unk_token_id @@ -85,8 +93,11 @@ def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torc mask = attention_mask.unsqueeze(-1).to(self.embeddings.weight.dtype) if self.unk_token_id is not None: mask[input_ids == self.unk_token_id] = 0 + embedding_ids = input_ids if self.token_mapping is None else self.token_mapping[input_ids] # Zero out padding - embeddings = self.embeddings(input_ids) * mask + embeddings = self.embeddings(embedding_ids) * mask + if self.weights is not None: + embeddings = embeddings * self.weights[input_ids].unsqueeze(-1) embeddings = embeddings.sum(dim=1) / mask.sum(dim=1).clamp(min=1) # Normalize if required if self.normalize: From 5c14d5381c0ac02b4e7bcb40b0654e63d4810793 Mon Sep 17 00:00:00 2001 From: serhiizghama Date: Fri, 11 Sep 2026 20:27:19 +0700 Subject: [PATCH 2/2] test: cover onnx export of vocabulary-quantized encoders and pipelines --- tests/test_export_to_onnx.py | 48 +++++++++++++++++++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/tests/test_export_to_onnx.py b/tests/test_export_to_onnx.py index 4feed2d..b229d24 100644 --- a/tests/test_export_to_onnx.py +++ b/tests/test_export_to_onnx.py @@ -17,7 +17,7 @@ from model2vec import StaticModel from model2vec.inference import StaticModelPipeline from model2vec.inference.mlp import Activation -from model2vec.model import DEFAULT_MAX_LENGTH +from model2vec.model import DEFAULT_MAX_LENGTH, quantize_model from model2vec.onnx import ( TorchStaticModel, TorchStaticModelPipeline, @@ -27,6 +27,7 @@ _save_tokenizer_and_config, export_model_to_onnx, ) +from model2vec.train import StaticModelForClassification def _tokenize(pipeline: StaticModelPipeline, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]: @@ -290,3 +291,48 @@ def test_encoder_onnx_without_unk_token(tmp_path: Path) -> None: expected = model.encode(texts) np.testing.assert_allclose(onnx_output, expected, atol=1e-5) + + +def _vocabulary_quantized_model() -> StaticModel: + """Build a small model whose vocabulary has been quantized, so ids go through `token_mapping` and `weights`.""" + vocab = ["[PAD]", "dog", "cat", "fish", "bird", "cow", "horse", "[UNK]"] + tokenizer = Tokenizer( + BPE(vocab={t: i for i, t in enumerate(vocab)}, merges=[], unk_token="[UNK]", ignore_merges=True) + ) + tokenizer.pre_tokenizer = Whitespace() # type: ignore[assignment] + vectors = np.random.RandomState(0).randn(len(vocab), 8).astype(np.float32) + model = quantize_model( + StaticModel(vectors=vectors, tokenizer=tokenizer, normalize=False), vocabulary_quantization=3 + ) + assert model.token_mapping is not None and model.weights is not None + assert len(model.embedding) < len(model.tokens) + return model + + +def test_encoder_onnx_vocabulary_quantized(tmp_path: Path) -> None: + """A vocabulary-quantized encoder maps ids through `token_mapping` and applies `weights`, like `encode`.""" + model = _vocabulary_quantized_model() + + texts = ["dog cat", "horse cow fish", "bird zzz"] + onnx_output = _encoder_onnx_output(model, texts, tmp_path / "model.onnx") + expected = model.encode(texts) + + np.testing.assert_allclose(onnx_output, expected, atol=1e-5) + + +def test_pipeline_onnx_vocabulary_quantized(tmp_path: Path) -> None: + """A classifier trained on a vocabulary-quantized model exports with its mapping and weights intact.""" + torch.random.manual_seed(42) + classifier = StaticModelForClassification.from_static_model(model=_vocabulary_quantized_model(), hidden_dim=8) + classifier.fit(["dog cat", "horse cow"], ["a", "b"]) + pipeline = classifier.to_pipeline() + assert pipeline.model.token_mapping is not None + + texts = ["dog cat", "horse cow fish", "bird"] + torch_model = TorchStaticModelPipeline(pipeline) + input_ids, attention_mask = _tokenize(pipeline, texts) + + onnx_output = _export(torch_model, input_ids, attention_mask, tmp_path / "model.onnx") + expected = pipeline.predict_proba(texts, use_multiprocessing=False) + + np.testing.assert_allclose(onnx_output, expected, atol=1e-4)