fix: keep float32 output for mean pooling models

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

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):

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):

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():

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(