diff --git a/README.md b/README.md index d4c88261..206915a0 100644 --- a/README.md +++ b/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 diff --git a/fastembed/common/utils.py b/fastembed/common/utils.py index 0e271fdb..fefa3238 100644 --- a/fastembed/common/utils.py +++ b/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 diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index c3d5cf5f..40d8e0ec 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/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 diff --git a/fastembed/text/pooled_embedding.py b/fastembed/text/pooled_embedding.py index 8e7b65c9..3dac5825 100644 --- a/fastembed/text/pooled_embedding.py +++ b/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): diff --git a/fastembed/text/pooled_normalized_embedding.py b/fastembed/text/pooled_normalized_embedding.py index 31bf4337..f7472857 100644 --- a/fastembed/text/pooled_normalized_embedding.py +++ b/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): diff --git a/tests/test_common.py b/tests/test_common.py index fe9626eb..ed2e087f 100644 --- a/tests/test_common.py +++ b/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]] diff --git a/tests/test_custom_models.py b/tests/test_custom_models.py index 1cf47f7f..33de31e7 100644 --- a/tests/test_custom_models.py +++ b/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(): diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 82e17c09..538c3335 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/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: