fix: keep the model dtype in mean pooling embeddings

README.md

@@ -169,10 +169,10 @@ query = "What is Qdrant?"
 model = LateInteractionMultimodalEmbedding(model_name="Qdrant/colpali-v1.3-fp16")
 doc_images_embeddings = list(model.embed_image(doc_images))
 # shape (2, 1030, 128)
-# [array([[-0.03353882, -0.02090454, ..., -0.15576172, -0.07678223]], dtype=float32)]
+# [array([[-0.03354, -0.0209, ..., -0.1558, -0.0768]], dtype=float16)]
 query_embedding = model.embed_text(query)
 # shape (1, 20, 128)
-# [array([[-0.00218201,  0.14758301, ...,  -0.02207947,  0.16833496]], dtype=float32)]
+# [array([[-0.002182, 0.1476, ..., -0.02208, 0.1683]], dtype=float16)]
 ```
 
 ### 🔄 Rerankers

fastembed/common/utils.py

@@ -16,6 +16,11 @@
 
 
 def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e-12) -> NumpyArray:
+    if input_array.dtype == np.float16:
+        # the sum of squares overflows float16 (max 65504) already for moderate values,
+        # which turns the norm into inf and the embedding into zeros
+        return normalize(input_array.astype(np.float32), p=p, dim=dim, eps=eps).astype(np.float16)
+
     # Calculate the Lp norm along the specified dimension
     norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
     norm = np.maximum(norm, eps)  # Avoid division by zero
@@ -24,10 +29,16 @@ def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e
 
 
 def mean_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) -> NumpyArray:
-    input_mask_expanded = np.expand_dims(attention_mask, axis=-1).astype(np.int64)
-    input_mask_expanded = np.tile(input_mask_expanded, (1, 1, input_array.shape[-1]))
-    sum_embeddings = np.sum(input_array * input_mask_expanded, axis=1)
-    sum_mask = np.sum(input_mask_expanded, axis=1)
+    """Average the embeddings of the tokens which the attention mask marks as real.
+
+    The sum is accumulated in float64, so the result is float64 for any input dtype,
+    callers cast it back to the dtype of the model once post-processing is done.
+    """
+    # `where` skips the padding without materializing a (batch_size, seq_len, dim) mask
+    sum_embeddings = np.sum(
+        input_array, axis=1, where=attention_mask[:, :, np.newaxis].astype(bool), dtype=np.float64
+    )
+    sum_mask = np.sum(attention_mask, axis=1, keepdims=True)
     pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
     return pooled_embeddings
 

fastembed/text/custom_text_embedding.py

@@ -72,7 +72,12 @@ 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))
+        # mean pooling returns float64, float embeddings are cast back to the dtype of the model
+        # after normalization, integer outputs are kept as is, since the cast would truncate them
+        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,8 @@ def _post_process_onnx_output(
 
         embeddings = output.model_output
         attn_mask = output.attention_mask
-        return self.mean_pooling(embeddings, attn_mask)
+        # mean pooling returns float64, embeddings keep the dtype of the model
+        return self.mean_pooling(embeddings, attn_mask).astype(embeddings.dtype, copy=False)
 
 
 class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):

fastembed/text/pooled_normalized_embedding.py

@@ -158,7 +158,11 @@ def _post_process_onnx_output(
 
         embeddings = output.model_output
         attn_mask = output.attention_mask
-        return normalize(self.mean_pooling(embeddings, attn_mask))
+        # mean pooling returns float64, embeddings keep the dtype of the model,
+        # the cast goes after normalization to normalize in full precision
+        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(
+            embeddings.dtype, copy=False
+        )
 
 
 class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):

tests/test_common.py

@@ -8,7 +8,7 @@
     LateInteractionMultimodalEmbedding,
     LateInteractionTextEmbedding,
 )
-from fastembed.common.utils import iter_batch, last_token_pooling
+from fastembed.common.utils import iter_batch, last_token_pooling, mean_pooling, normalize
 
 
 def test_text_list_supported_models():
@@ -62,6 +62,31 @@ def test_last_token_pooling_with_left_padding():
     assert np.allclose(pooled, [[2.0, 2.0], [6.0, 6.0]])
 
 
+def test_mean_pooling():
+    # 2 real tokens, then padding
+    token_embeddings = np.array([[[1.0, 2.0], [3.0, 4.0], [9.0, 9.0]]], dtype=np.float32)
+    attention_mask = np.array([[1, 1, 0]], dtype=np.int64)
+    # the sum over 8192 tokens of 10.0 exceeds the float16 max of 65504
+    long_sequence = np.full((1, 8192, 2), 10.0, dtype=np.float16)
+
+    pooled = mean_pooling(token_embeddings, attention_mask)
+    pooled_long_sequence = mean_pooling(long_sequence, np.ones((1, 8192), dtype=np.int64))
+
+    assert pooled.dtype == pooled_long_sequence.dtype == np.float64
+    assert np.array_equal(pooled, [[2.0, 3.0]])
+    assert np.array_equal(pooled_long_sequence, [[10.0, 10.0]])
+
+
+def test_normalize_does_not_overflow_float16():
+    # the sum of squares, 1024 * 10.0**2, exceeds the float16 max of 65504
+    embeddings = np.full((1, 1024), 10.0, dtype=np.float16)
+
+    normalized = normalize(embeddings)
+
+    assert normalized.dtype == np.float16
+    assert np.array_equal(normalized, np.full((1, 1024), 1 / 32))
+
+
 def test_iter_batch_accepts_positive_size():
     assert list(iter_batch([1, 2, 3, 4, 5], 3)) == [[1, 2, 3], [4, 5]]
 

tests/test_custom_models.py

@@ -225,6 +225,7 @@ 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, model_name
 
 
 def test_custom_text_model_lookup_is_case_insensitive():

tests/test_text_onnx_embeddings.py

@@ -5,6 +5,8 @@
 import numpy as np
 import pytest
 
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.text.custom_text_embedding import CustomTextEmbedding
 from fastembed.text.last_token_normalized_embedding import LastTokenNormalizedEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding import TextEmbedding
@@ -241,6 +243,30 @@ def test_quantized_model_reports_onnxruntime_requirement(monkeypatch) -> None:
         model.load_onnx_model()
 
 
+@pytest.mark.parametrize(
+    "embedding_class",
+    # custom models are covered in test_custom_models.py
+    [cls for cls in TextEmbedding.EMBEDDINGS_REGISTRY if cls is not CustomTextEmbedding],
+    ids=lambda cls: cls.__name__,
+)
+def test_post_processing_keeps_model_dtype(embedding_class) -> None:
+    model_desc = embedding_class._list_supported_models()[0]
+    model = embedding_class(
+        model_desc.model,
+        lazy_load=True,
+        specific_model_path="./",  # disable model downloading and loading
+    )
+    token_embeddings = np.random.default_rng(0).standard_normal((2, 4, model_desc.dim))
+    output = OnnxOutputContext(
+        model_output=token_embeddings.astype(np.float32),
+        attention_mask=np.array([[1, 1, 0, 0], [1, 1, 1, 1]], dtype=np.int64),
+    )
+
+    embeddings = np.stack(list(model._post_process_onnx_output(output)))
+
+    assert embeddings.dtype == np.float32
+
+
 @pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
 def test_batch_embedding(model_cache, n_dims: int, model_name: str) -> None:
     with model_cache(model_name) as model: