diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index c3d5cf5f..caa6872e 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/fastembed/text/custom_text_embedding.py @@ -72,7 +72,17 @@ def _get_worker_init_kwargs(self) -> dict[str, Any]: def _post_process_onnx_output( self, output: OnnxOutputContext, **kwargs: Any ) -> Iterable[NumpyArray]: - return self._normalize(self._pool(output.model_output, output.attention_mask)) + """Pool and optionally normalize the model output, as configured for this custom model. + + Float embeddings keep the dtype the model produced, so a float32 model returns float32 + even though mean pooling computes in float64. Integer outputs are returned as computed, + since casting them back would truncate the pooled values. + """ + embeddings = self._normalize(self._pool(output.model_output, output.attention_mask)) + # cast after normalizing: computing the norm in float16 can overflow + if np.issubdtype(output.model_output.dtype, np.floating): + return embeddings.astype(output.model_output.dtype, copy=False) + return embeddings def _pool( self, embeddings: NumpyArray, attention_mask: NDArray[np.int64] | None = None diff --git a/fastembed/text/pooled_embedding.py b/fastembed/text/pooled_embedding.py index 8e7b65c9..c140e067 100644 --- a/fastembed/text/pooled_embedding.py +++ b/fastembed/text/pooled_embedding.py @@ -117,7 +117,7 @@ def _post_process_onnx_output( embeddings = output.model_output attn_mask = output.attention_mask - return self.mean_pooling(embeddings, attn_mask) + return self.mean_pooling(embeddings, attn_mask).astype(embeddings.dtype, copy=False) class PooledEmbeddingWorker(OnnxTextEmbeddingWorker): diff --git a/fastembed/text/pooled_normalized_embedding.py b/fastembed/text/pooled_normalized_embedding.py index 31bf4337..72ddf809 100644 --- a/fastembed/text/pooled_normalized_embedding.py +++ b/fastembed/text/pooled_normalized_embedding.py @@ -158,7 +158,9 @@ def _post_process_onnx_output( embeddings = output.model_output attn_mask = output.attention_mask - return normalize(self.mean_pooling(embeddings, attn_mask)) + return normalize(self.mean_pooling(embeddings, attn_mask)).astype( + embeddings.dtype, copy=False + ) class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker): diff --git a/tests/test_custom_models.py b/tests/test_custom_models.py index 1cf47f7f..8ef846e5 100644 --- a/tests/test_custom_models.py +++ b/tests/test_custom_models.py @@ -225,6 +225,35 @@ def test_mock_add_custom_models(): iter(custom_text_embedding._post_process_onnx_output(input_data[model_name])) ) assert np.allclose(post_processed_output, expected_output[model_name], atol=1e-3) + assert post_processed_output.dtype == np.float32 + + +def test_custom_mean_pooling_keeps_float16_without_overflow(): + """A float16 model stays float16, and normalizing before the cast keeps the norm finite.""" + model_name = "mean-normalized-fp16" + TextEmbedding.add_custom_model( + model_name, + pooling=PoolingType.MEAN, + normalization=True, + sources=ModelSource(hf="artificial"), + dim=1024, + size_in_gb=0.1, + ) + custom_text_embedding = CustomTextEmbedding( + model_name, + lazy_load=True, + specific_model_path="./", # disable model downloading and loading + ) + # the squared norm, 1024 * 10.0**2, exceeds the float16 max of 65504 + output = OnnxOutputContext( + model_output=np.full((1, 3, 1024), 10.0, dtype=np.float16), + attention_mask=np.ones((1, 3), dtype=np.int64), + ) + + embedding = next(iter(custom_text_embedding._post_process_onnx_output(output))) + + assert embedding.dtype == np.float16 + assert np.allclose(embedding, 1 / 32, atol=1e-3) def test_custom_text_model_lookup_is_case_insensitive(): diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 82e17c09..f1fc2b91 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -5,8 +5,12 @@ import numpy as np import pytest +from fastembed.common.onnx_model import OnnxOutputContext +from fastembed.common.utils import mean_pooling, normalize from fastembed.text.last_token_normalized_embedding import LastTokenNormalizedEmbedding from fastembed.text.onnx_embedding import OnnxTextEmbedding +from fastembed.text.pooled_embedding import PooledEmbedding +from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding from fastembed.text.text_embedding import TextEmbedding from tests.utils import delete_model_cache, should_test_model @@ -182,6 +186,7 @@ def test_embedding(model_cache, model_name: str) -> None: docs = [DOC_PREFIXES[model_desc.model] + doc for doc in docs] embeddings = list(model.embed(docs)) + assert embeddings[0].dtype == np.float32, model_desc.model embeddings = np.stack(embeddings, axis=0) assert embeddings.shape == (2, dim) @@ -224,6 +229,37 @@ def test_query_embedding(model_cache) -> None: ), model_desc.model +@pytest.mark.parametrize( + "model_class,model_name,normalized", + [ + (PooledEmbedding, "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", False), + (PooledNormalizedEmbedding, "sentence-transformers/all-MiniLM-L6-v2", True), + ], +) +@pytest.mark.parametrize("dtype", [np.float32, np.float16]) +def test_mean_pooling_keeps_model_dtype( + model_class, model_name: str, normalized: bool, dtype +) -> None: + """Built-in mean pooling returns the model's own dtype and the correct pooled values.""" + model = model_class( + model_name, + lazy_load=True, + specific_model_path="./", # disable model downloading and loading + ) + token_embeddings = np.random.random((2, 4, 8)).astype(dtype) + attention_mask = np.array([[1, 1, 1, 1], [1, 1, 0, 0]], dtype=np.int64) + output = OnnxOutputContext(model_output=token_embeddings, attention_mask=attention_mask) + + expected = mean_pooling(token_embeddings, attention_mask) + if normalized: + expected = normalize(expected) + + embeddings = model._post_process_onnx_output(output) + + assert embeddings.dtype == dtype + assert np.allclose(embeddings, expected, atol=1e-3) + + def test_quantized_model_reports_onnxruntime_requirement(monkeypatch) -> None: """Old onnxruntime only implements 4-bit MatMulNBits, the error should say so.""" monkeypatch.setattr(