From f4811cc42d4d43725c6416dcabf687e1744d8d32 Mon Sep 17 00:00:00 2001 From: stephantul Date: Mon, 14 Sep 2026 21:28:46 +0200 Subject: [PATCH 1/5] feat: make encoding faster --- model2vec/model.py | 34 +++++++++++++++++++++------------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/model2vec/model.py b/model2vec/model.py index 6c7dde6..65e9aaa 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -4,6 +4,7 @@ import json import math import os +from collections import defaultdict from collections.abc import Iterator, Sequence from logging import getLogger from pathlib import Path @@ -385,8 +386,8 @@ def encode( show_progress_bar: bool = False, max_length: int | None | _UnsetType = _UNSET, normalize: bool | None = None, - batch_size: int = 1024, - use_multiprocessing: bool = True, + batch_size: int = 131072, + use_multiprocessing: bool = False, multiprocessing_threshold: int = 10_000, **kwargs: Any, ) -> np.ndarray: @@ -455,20 +456,27 @@ def _encode_helper(self, id_list: list[int]) -> np.ndarray: def _encode_batch(self, sentences: Sequence[str], normalize: bool) -> np.ndarray: """Encode a batch of sentences.""" ids = self.tokenize(sentences=sentences) - out: list[np.ndarray] = [] - for id_list in ids: - if id_list: - emb = self._encode_helper(id_list) - out.append(emb.mean(axis=0)) - else: - out.append(np.zeros(self.dim)) + out = np.zeros((len(ids), self.dim), dtype=self.embedding_dtype) + + if self.token_mapping is None and self.weights is None: + buckets: dict[int, list[int]] = defaultdict(list) + for i, id_list in enumerate(ids): + if id_list: + buckets[len(id_list)].append(i) + for _, indices in buckets.items(): + id_matrix = np.array([ids[i] for i in indices], dtype=np.int64) + out[indices] = self.embedding[id_matrix].mean(axis=1) + else: + for i, id_list in enumerate(ids): + if id_list: + emb = self._encode_helper(id_list) + out[i] = emb.mean(axis=0) - out_array = np.stack(out) if normalize: - norm = np.linalg.norm(out_array, axis=1, keepdims=True) + 1e-32 - out_array = out_array / norm + norm = np.linalg.norm(out, axis=1, keepdims=True) + 1e-32 + np.divide(out, norm, out=out) - return out_array + return out @staticmethod def _batch(sentences: Sequence[str], batch_size: int) -> Iterator[Sequence[str]]: From d2883efc4e21aff39dd5d9fffebe39065b8bb6b0 Mon Sep 17 00:00:00 2001 From: stephantul Date: Tue, 15 Sep 2026 08:56:01 +0200 Subject: [PATCH 2/5] threading-based paralellism --- model2vec/model.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model2vec/model.py b/model2vec/model.py index 65e9aaa..057ab97 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -280,9 +280,9 @@ def _encode_dispatch( # Disable parallelism for tokenizers os.environ["TOKENIZERS_PARALLELISM"] = "false" - results = ProgressParallel(n_jobs=-1, use_tqdm=show_progress_bar, total=total_batches)( - delayed(batch_fn)(batch, *batch_args) for batch in sentence_batches - ) + results = ProgressParallel( + n_jobs=-1, backend="threading", use_tqdm=show_progress_bar, total=total_batches + )(delayed(batch_fn)(batch, *batch_args) for batch in sentence_batches) else: results = [ batch_fn(batch, *batch_args) From 7d08ad75351ddf52784e5b896c2a0411c4adbf7a Mon Sep 17 00:00:00 2001 From: stephantul Date: Tue, 15 Sep 2026 10:14:08 +0200 Subject: [PATCH 3/5] turn on mp again --- model2vec/model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/model2vec/model.py b/model2vec/model.py index 057ab97..1c693b4 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -386,8 +386,8 @@ def encode( show_progress_bar: bool = False, max_length: int | None | _UnsetType = _UNSET, normalize: bool | None = None, - batch_size: int = 131072, - use_multiprocessing: bool = False, + batch_size: int = 1024, + use_multiprocessing: bool = True, multiprocessing_threshold: int = 10_000, **kwargs: Any, ) -> np.ndarray: From b81facb67a13181e61ab0e994d0858ba638801e6 Mon Sep 17 00:00:00 2001 From: stephantul Date: Tue, 15 Sep 2026 10:17:33 +0200 Subject: [PATCH 4/5] proactively cast to float32 for int8 models --- model2vec/model.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/model2vec/model.py b/model2vec/model.py index 1c693b4..324062c 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -456,7 +456,10 @@ def _encode_helper(self, id_list: list[int]) -> np.ndarray: def _encode_batch(self, sentences: Sequence[str], normalize: bool) -> np.ndarray: """Encode a batch of sentences.""" ids = self.tokenize(sentences=sentences) - out = np.zeros((len(ids), self.dim), dtype=self.embedding_dtype) + dtype = self.embedding.dtype + if dtype == np.int8: + dtype = np.float32 + out = np.zeros((len(ids), self.dim), dtype=dtype) if self.token_mapping is None and self.weights is None: buckets: dict[int, list[int]] = defaultdict(list) From 85ad9fe4758c46a20278076b56737cd1048b9cc2 Mon Sep 17 00:00:00 2001 From: stephantul Date: Tue, 15 Sep 2026 11:33:47 +0200 Subject: [PATCH 5/5] add test --- tests/test_model.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/test_model.py b/tests/test_model.py index 66c114d..df9f25e 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -62,6 +62,20 @@ def test_encode_multiple_sentences( assert encoded.shape == (2, 2) +def test_encode_int8_quantized( + mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str] +) -> None: + """Test that encoding an int8-quantized model returns float32 output.""" + from model2vec.model import quantize_model + + model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config) + quantized = quantize_model(model, quantize_to="int8") + assert quantized.embedding.dtype == np.int8 + + encoded = quantized.encode(["word1 word2", "word1 word3"]) + assert encoded.dtype == np.float32 + + def test_encode_as_sequence(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str]) -> None: """Test encoding of sentences as tokens.""" sentences = ["word1 word2", "word1 word3"]