From 7c705c937df1399f6d519c9852b9051b6a05d9af Mon Sep 17 00:00:00 2001 From: jaysheel-dodia Date: Tue, 29 Sep 2026 21:39:10 -0400 Subject: [PATCH 1/3] fix: keep model dtype in mean pooling --- fastembed/text/custom_text_embedding.py | 6 +++- fastembed/text/pooled_embedding.py | 2 +- fastembed/text/pooled_normalized_embedding.py | 2 +- tests/test_custom_models.py | 28 ++++++++++++++++ tests/test_text_onnx_embeddings.py | 33 +++++++++++++++++++ 5 files changed, 68 insertions(+), 3 deletions(-) diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index c3d5cf5f0..920ba415c 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/fastembed/text/custom_text_embedding.py @@ -72,7 +72,11 @@ 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)) + 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 8e7b65c9a..c140e0671 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 31bf43377..68bf40e39 100644 --- a/fastembed/text/pooled_normalized_embedding.py +++ b/fastembed/text/pooled_normalized_embedding.py @@ -158,7 +158,7 @@ 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 1cf47f7f1..9e4b1ff27 100644 --- a/tests/test_custom_models.py +++ b/tests/test_custom_models.py @@ -225,6 +225,34 @@ 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(): + 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 82e17c09d..42c416f7a 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,34 @@ 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: + 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( From 4b8dd8aae94d4360671c4954d6be4d40005e2a0d Mon Sep 17 00:00:00 2001 From: jaysheel-dodia Date: Tue, 29 Sep 2026 22:08:17 -0400 Subject: [PATCH 2/3] style: format long lines --- fastembed/text/pooled_normalized_embedding.py | 4 +++- tests/test_text_onnx_embeddings.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/fastembed/text/pooled_normalized_embedding.py b/fastembed/text/pooled_normalized_embedding.py index 68bf40e39..72ddf8095 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)).astype(embeddings.dtype, copy=False) + return normalize(self.mean_pooling(embeddings, attn_mask)).astype( + embeddings.dtype, copy=False + ) class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker): diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 42c416f7a..a61c3e6ad 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -237,7 +237,9 @@ def test_query_embedding(model_cache) -> None: ], ) @pytest.mark.parametrize("dtype", [np.float32, np.float16]) -def test_mean_pooling_keeps_model_dtype(model_class, model_name: str, normalized: bool, dtype) -> None: +def test_mean_pooling_keeps_model_dtype( + model_class, model_name: str, normalized: bool, dtype +) -> None: model = model_class( model_name, lazy_load=True, From 2b752dbdd9a280eb9999320d8b0bc43829443ed2 Mon Sep 17 00:00:00 2001 From: jaysheel-dodia Date: Tue, 29 Sep 2026 22:41:11 -0400 Subject: [PATCH 3/3] docs: add docstrings for mean pooling dtype fix --- fastembed/text/custom_text_embedding.py | 6 ++++++ tests/test_custom_models.py | 1 + tests/test_text_onnx_embeddings.py | 1 + 3 files changed, 8 insertions(+) diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index 920ba415c..caa6872e5 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/fastembed/text/custom_text_embedding.py @@ -72,6 +72,12 @@ def _get_worker_init_kwargs(self) -> dict[str, Any]: def _post_process_onnx_output( self, output: OnnxOutputContext, **kwargs: Any ) -> Iterable[NumpyArray]: + """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): diff --git a/tests/test_custom_models.py b/tests/test_custom_models.py index 9e4b1ff27..8ef846e56 100644 --- a/tests/test_custom_models.py +++ b/tests/test_custom_models.py @@ -229,6 +229,7 @@ def test_mock_add_custom_models(): 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, diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index a61c3e6ad..f1fc2b91b 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -240,6 +240,7 @@ def test_query_embedding(model_cache) -> None: 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,