From 8f1bd82f9a1d1c280accd53dfb2d85b87e34810b Mon Sep 17 00:00:00 2001 From: George Date: Wed, 30 Sep 2026 16:30:57 +0700 Subject: [PATCH 1/2] fix: keep the model dtype in mean pooling embeddings --- README.md | 4 +- fastembed/common/utils.py | 19 ++++++++-- fastembed/text/custom_text_embedding.py | 7 +++- fastembed/text/pooled_embedding.py | 3 +- fastembed/text/pooled_normalized_embedding.py | 6 ++- tests/test_common.py | 37 ++++++++++++++++++- tests/test_custom_models.py | 34 +++++++++++++++-- tests/test_image_onnx_embeddings.py | 23 ++++++++++++ tests/test_text_onnx_embeddings.py | 32 ++++++++++++++++ 9 files changed, 152 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index d4c882615..206915a03 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 0e271fdb3..fefa3238b 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 c3d5cf5f0..40d8e0eca 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 8e7b65c9a..3dac58256 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 31bf43377..f74728578 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 fe9626ebf..80b0c4baf 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,41 @@ def test_last_token_pooling_with_left_padding(): assert np.allclose(pooled, [[2.0, 2.0], [6.0, 6.0]]) +@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64]) +def test_mean_pooling(dtype): + token_embeddings = np.array( + [ + [[1.0, 2.0], [3.0, 4.0], [9.0, 9.0]], # 2 real tokens, then padding + [[9.0, 9.0], [1.0, 2.0], [5.0, 8.0]], # padding, then 2 real tokens + [[9.0, 9.0], [9.0, 9.0], [9.0, 9.0]], # padding only + ], + dtype=dtype, + ) + attention_mask = np.array([[1, 1, 0], [0, 1, 1], [0, 0, 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=dtype) + + 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], [3.0, 5.0], [0.0, 0.0]]) + assert np.array_equal(pooled_long_sequence, [[10.0, 10.0]]) + + +@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64]) +def test_normalize(dtype): + # the sum of squares of the first row, 1024 * 10.0**2, exceeds the float16 max of 65504 + embeddings = np.zeros((2, 1024), dtype=dtype) + embeddings[0] = 10.0 + + normalized = normalize(embeddings) + + assert normalized.dtype == dtype + assert np.array_equal(normalized[0], np.full(1024, 1 / 32)) + assert not normalized[1].any() # a zero vector stays zero + + 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 1cf47f7f1..cf4303ed1 100644 --- a/tests/test_custom_models.py +++ b/tests/test_custom_models.py @@ -70,6 +70,7 @@ def test_text_custom_model(): embeddings = list(model.embed(docs)) embeddings = np.stack(embeddings, axis=0) assert embeddings.shape == (2, dim) + assert embeddings.dtype == np.float32 assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3) @@ -96,6 +97,7 @@ def test_text_custom_model_parallel_processing(): embeddings = np.stack(list(model.embed(docs, batch_size=10, parallel=2)), axis=0) assert embeddings.shape == (len(docs), dim) + assert embeddings.dtype == np.float32 if is_ci: delete_model_cache(model.model._model_dir) @@ -159,14 +161,15 @@ def test_cross_encoder_custom_model_parallel_processing(): delete_model_cache(model.model._model_dir) -def test_mock_add_custom_models(): +@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64]) +def test_mock_add_custom_models(dtype): dim = 5 size_in_gb = 0.1 source = ModelSource(hf="artificial") num_tokens = 10 - dummy_pooled_embedding = np.random.random((1, dim)).astype(np.float32) - dummy_token_embedding = np.random.random((1, num_tokens, dim)).astype(np.float32) + dummy_pooled_embedding = np.random.random((1, dim)).astype(dtype) + dummy_token_embedding = np.random.random((1, num_tokens, dim)).astype(dtype) dummy_attention_mask = np.ones((1, num_tokens)).astype(np.int64) dummy_token_output = OnnxOutputContext( @@ -225,6 +228,31 @@ 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 == dtype, model_name + + +def test_mock_custom_mean_pooling_does_not_truncate_integer_outputs() -> None: + TextEmbedding.add_custom_model( + "mean-int8", + pooling=PoolingType.MEAN, + normalization=False, + sources=ModelSource(hf="artificial"), + dim=2, + size_in_gb=0.1, + ) + custom_text_embedding = CustomTextEmbedding( + "mean-int8", + lazy_load=True, + specific_model_path="./", # disable model downloading and loading + ) + output = OnnxOutputContext( + model_output=np.array([[[1, 2], [2, 3]]], dtype=np.int8), + attention_mask=np.ones((1, 2), dtype=np.int64), + ) + + embedding = next(iter(custom_text_embedding._post_process_onnx_output(output))) + + assert np.array_equal(embedding, [1.5, 2.5]) def test_custom_text_model_lookup_is_case_insensitive(): diff --git a/tests/test_image_onnx_embeddings.py b/tests/test_image_onnx_embeddings.py index 7774055bb..898e7a8b1 100644 --- a/tests/test_image_onnx_embeddings.py +++ b/tests/test_image_onnx_embeddings.py @@ -9,6 +9,7 @@ from PIL import Image from fastembed import ImageEmbedding +from fastembed.common.onnx_model import OnnxOutputContext from tests.config import TEST_MISC_DIR from tests.utils import delete_model_cache, should_test_model @@ -91,6 +92,7 @@ def test_embedding(model_cache, model_name: str) -> None: embeddings = list(model.embed(images)) embeddings = np.stack(embeddings, axis=0) assert embeddings.shape == (len(images), dim) + assert embeddings.dtype == np.float32, model_desc.model canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model] @@ -161,6 +163,27 @@ def test_lazy_load(model_name: str) -> None: delete_model_cache(model.model._model_dir) +@pytest.mark.parametrize("dtype", [np.float32, np.float16]) +@pytest.mark.parametrize( + "embedding_class", ImageEmbedding.EMBEDDINGS_REGISTRY, ids=lambda cls: cls.__name__ +) +def test_post_processing_keeps_model_dtype(embedding_class, dtype) -> 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 + ) + # values this large overflow the sum of squares in float16, unless it's computed in a wider type + image_embeddings = np.random.default_rng(0).standard_normal((2, model_desc.dim)) * 100 + output = OnnxOutputContext(model_output=image_embeddings.astype(dtype)) + + embeddings = np.stack(list(model._post_process_onnx_output(output))) + + assert embeddings.dtype == dtype + assert np.isfinite(embeddings).all() and embeddings.any() + + def test_get_embedding_size() -> None: assert ImageEmbedding.get_embedding_size(model_name="Qdrant/clip-ViT-B-32-vision") == 512 assert ImageEmbedding.get_embedding_size(model_name="Qdrant/clip-vit-b-32-vision") == 512 diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 82e17c09d..9ffacc631 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 @@ -184,6 +186,7 @@ def test_embedding(model_cache, model_name: str) -> None: embeddings = list(model.embed(docs)) embeddings = np.stack(embeddings, axis=0) assert embeddings.shape == (2, dim) + assert embeddings.dtype == np.float32, model_desc.model canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model] assert np.allclose( @@ -217,6 +220,7 @@ def test_query_embedding(model_cache) -> None: embeddings = list(model.query_embed(queries)) embeddings = np.stack(embeddings, axis=0) assert embeddings.shape == (2, dim) + assert embeddings.dtype == np.float32, model_desc.model canonical_vector = CANONICAL_QUERY_VECTOR_VALUES[model_desc.model] assert np.allclose( @@ -241,6 +245,33 @@ def test_quantized_model_reports_onnxruntime_requirement(monkeypatch) -> None: model.load_onnx_model() +@pytest.mark.parametrize("dtype", [np.float32, np.float16]) +@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, dtype) -> 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 + ) + # values this large overflow the sum of squares in float16, unless it's computed in a wider type + token_embeddings = np.random.default_rng(0).standard_normal((2, 4, model_desc.dim)) * 100 + output = OnnxOutputContext( + model_output=token_embeddings.astype(dtype), + 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 == dtype + assert np.isfinite(embeddings).all() and embeddings.any() + + @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: @@ -265,6 +296,7 @@ def test_parallel_processing(model_cache, n_dims: int, model_name: str) -> None: embeddings_3 = np.stack(embeddings_3, axis=0) assert embeddings.shape == (len(docs), n_dims) + assert embeddings.dtype == np.float32 assert np.allclose(embeddings, embeddings_2, atol=1e-3) assert np.allclose(embeddings, embeddings_3, atol=1e-3) From c1cdadb9a716aad82f0e7373a798faf06751e05e Mon Sep 17 00:00:00 2001 From: George Date: Wed, 30 Sep 2026 18:27:04 +0700 Subject: [PATCH 2/2] tests: trim redundant dtype tests --- tests/test_common.py | 32 +++++++++----------------- tests/test_custom_models.py | 35 ++++------------------------- tests/test_image_onnx_embeddings.py | 23 ------------------- tests/test_text_onnx_embeddings.py | 14 ++++-------- 4 files changed, 19 insertions(+), 85 deletions(-) diff --git a/tests/test_common.py b/tests/test_common.py index 80b0c4baf..ed2e087ff 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -62,39 +62,29 @@ def test_last_token_pooling_with_left_padding(): assert np.allclose(pooled, [[2.0, 2.0], [6.0, 6.0]]) -@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64]) -def test_mean_pooling(dtype): - token_embeddings = np.array( - [ - [[1.0, 2.0], [3.0, 4.0], [9.0, 9.0]], # 2 real tokens, then padding - [[9.0, 9.0], [1.0, 2.0], [5.0, 8.0]], # padding, then 2 real tokens - [[9.0, 9.0], [9.0, 9.0], [9.0, 9.0]], # padding only - ], - dtype=dtype, - ) - attention_mask = np.array([[1, 1, 0], [0, 1, 1], [0, 0, 0]], dtype=np.int64) +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=dtype) + 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], [3.0, 5.0], [0.0, 0.0]]) + assert np.array_equal(pooled, [[2.0, 3.0]]) assert np.array_equal(pooled_long_sequence, [[10.0, 10.0]]) -@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64]) -def test_normalize(dtype): - # the sum of squares of the first row, 1024 * 10.0**2, exceeds the float16 max of 65504 - embeddings = np.zeros((2, 1024), dtype=dtype) - embeddings[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 == dtype - assert np.array_equal(normalized[0], np.full(1024, 1 / 32)) - assert not normalized[1].any() # a zero vector stays zero + assert normalized.dtype == np.float16 + assert np.array_equal(normalized, np.full((1, 1024), 1 / 32)) def test_iter_batch_accepts_positive_size(): diff --git a/tests/test_custom_models.py b/tests/test_custom_models.py index cf4303ed1..33de31e79 100644 --- a/tests/test_custom_models.py +++ b/tests/test_custom_models.py @@ -70,7 +70,6 @@ def test_text_custom_model(): embeddings = list(model.embed(docs)) embeddings = np.stack(embeddings, axis=0) assert embeddings.shape == (2, dim) - assert embeddings.dtype == np.float32 assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3) @@ -97,7 +96,6 @@ def test_text_custom_model_parallel_processing(): embeddings = np.stack(list(model.embed(docs, batch_size=10, parallel=2)), axis=0) assert embeddings.shape == (len(docs), dim) - assert embeddings.dtype == np.float32 if is_ci: delete_model_cache(model.model._model_dir) @@ -161,15 +159,14 @@ def test_cross_encoder_custom_model_parallel_processing(): delete_model_cache(model.model._model_dir) -@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64]) -def test_mock_add_custom_models(dtype): +def test_mock_add_custom_models(): dim = 5 size_in_gb = 0.1 source = ModelSource(hf="artificial") num_tokens = 10 - dummy_pooled_embedding = np.random.random((1, dim)).astype(dtype) - dummy_token_embedding = np.random.random((1, num_tokens, dim)).astype(dtype) + dummy_pooled_embedding = np.random.random((1, dim)).astype(np.float32) + dummy_token_embedding = np.random.random((1, num_tokens, dim)).astype(np.float32) dummy_attention_mask = np.ones((1, num_tokens)).astype(np.int64) dummy_token_output = OnnxOutputContext( @@ -228,31 +225,7 @@ def test_mock_add_custom_models(dtype): 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 == dtype, model_name - - -def test_mock_custom_mean_pooling_does_not_truncate_integer_outputs() -> None: - TextEmbedding.add_custom_model( - "mean-int8", - pooling=PoolingType.MEAN, - normalization=False, - sources=ModelSource(hf="artificial"), - dim=2, - size_in_gb=0.1, - ) - custom_text_embedding = CustomTextEmbedding( - "mean-int8", - lazy_load=True, - specific_model_path="./", # disable model downloading and loading - ) - output = OnnxOutputContext( - model_output=np.array([[[1, 2], [2, 3]]], dtype=np.int8), - attention_mask=np.ones((1, 2), dtype=np.int64), - ) - - embedding = next(iter(custom_text_embedding._post_process_onnx_output(output))) - - assert np.array_equal(embedding, [1.5, 2.5]) + assert post_processed_output.dtype == np.float32, model_name def test_custom_text_model_lookup_is_case_insensitive(): diff --git a/tests/test_image_onnx_embeddings.py b/tests/test_image_onnx_embeddings.py index 898e7a8b1..7774055bb 100644 --- a/tests/test_image_onnx_embeddings.py +++ b/tests/test_image_onnx_embeddings.py @@ -9,7 +9,6 @@ from PIL import Image from fastembed import ImageEmbedding -from fastembed.common.onnx_model import OnnxOutputContext from tests.config import TEST_MISC_DIR from tests.utils import delete_model_cache, should_test_model @@ -92,7 +91,6 @@ def test_embedding(model_cache, model_name: str) -> None: embeddings = list(model.embed(images)) embeddings = np.stack(embeddings, axis=0) assert embeddings.shape == (len(images), dim) - assert embeddings.dtype == np.float32, model_desc.model canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model] @@ -163,27 +161,6 @@ def test_lazy_load(model_name: str) -> None: delete_model_cache(model.model._model_dir) -@pytest.mark.parametrize("dtype", [np.float32, np.float16]) -@pytest.mark.parametrize( - "embedding_class", ImageEmbedding.EMBEDDINGS_REGISTRY, ids=lambda cls: cls.__name__ -) -def test_post_processing_keeps_model_dtype(embedding_class, dtype) -> 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 - ) - # values this large overflow the sum of squares in float16, unless it's computed in a wider type - image_embeddings = np.random.default_rng(0).standard_normal((2, model_desc.dim)) * 100 - output = OnnxOutputContext(model_output=image_embeddings.astype(dtype)) - - embeddings = np.stack(list(model._post_process_onnx_output(output))) - - assert embeddings.dtype == dtype - assert np.isfinite(embeddings).all() and embeddings.any() - - def test_get_embedding_size() -> None: assert ImageEmbedding.get_embedding_size(model_name="Qdrant/clip-ViT-B-32-vision") == 512 assert ImageEmbedding.get_embedding_size(model_name="Qdrant/clip-vit-b-32-vision") == 512 diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 9ffacc631..538c33358 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -186,7 +186,6 @@ def test_embedding(model_cache, model_name: str) -> None: embeddings = list(model.embed(docs)) embeddings = np.stack(embeddings, axis=0) assert embeddings.shape == (2, dim) - assert embeddings.dtype == np.float32, model_desc.model canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model] assert np.allclose( @@ -220,7 +219,6 @@ def test_query_embedding(model_cache) -> None: embeddings = list(model.query_embed(queries)) embeddings = np.stack(embeddings, axis=0) assert embeddings.shape == (2, dim) - assert embeddings.dtype == np.float32, model_desc.model canonical_vector = CANONICAL_QUERY_VECTOR_VALUES[model_desc.model] assert np.allclose( @@ -245,31 +243,28 @@ def test_quantized_model_reports_onnxruntime_requirement(monkeypatch) -> None: model.load_onnx_model() -@pytest.mark.parametrize("dtype", [np.float32, np.float16]) @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, dtype) -> None: +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 ) - # values this large overflow the sum of squares in float16, unless it's computed in a wider type - token_embeddings = np.random.default_rng(0).standard_normal((2, 4, model_desc.dim)) * 100 + token_embeddings = np.random.default_rng(0).standard_normal((2, 4, model_desc.dim)) output = OnnxOutputContext( - model_output=token_embeddings.astype(dtype), + 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 == dtype - assert np.isfinite(embeddings).all() and embeddings.any() + assert embeddings.dtype == np.float32 @pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")]) @@ -296,7 +291,6 @@ def test_parallel_processing(model_cache, n_dims: int, model_name: str) -> None: embeddings_3 = np.stack(embeddings_3, axis=0) assert embeddings.shape == (len(docs), n_dims) - assert embeddings.dtype == np.float32 assert np.allclose(embeddings, embeddings_2, atol=1e-3) assert np.allclose(embeddings, embeddings_3, atol=1e-3)