diff --git a/model2vec/model.py b/model2vec/model.py index 6c7dde6..324062c 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 @@ -279,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) @@ -455,20 +456,30 @@ 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)) + 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) + 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]]: 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"]