From 3fc4ea38dc3c5996d24ffe26f2f762f7f59ee25e Mon Sep 17 00:00:00 2001 From: debaditya Date: Fri, 3 Jul 2026 13:03:20 +0530 Subject: [PATCH] fix: support quantization ranges for int8/uint8 sentence-transformers embeddings --- .../embedding_backend/backend.py | 9 +++- ...sentence_transformers_document_embedder.py | 22 ++++++++- .../sentence_transformers_text_embedder.py | 21 ++++++++- ...sentence_transformers_document_embedder.py | 14 ++++++ ...sentence_transformers_embedding_backend.py | 33 ++++++++++++++ ...est_sentence_transformers_text_embedder.py | 45 +++++++++++++++++++ 6 files changed, 141 insertions(+), 3 deletions(-) diff --git a/integrations/sentence_transformers/src/haystack_integrations/components/embedders/sentence_transformers/embedding_backend/backend.py b/integrations/sentence_transformers/src/haystack_integrations/components/embedders/sentence_transformers/embedding_backend/backend.py index 7c64155694..6c49b7a0c3 100644 --- a/integrations/sentence_transformers/src/haystack_integrations/components/embedders/sentence_transformers/embedding_backend/backend.py +++ b/integrations/sentence_transformers/src/haystack_integrations/components/embedders/sentence_transformers/embedding_backend/backend.py @@ -5,10 +5,11 @@ import json from typing import Any, ClassVar, Literal +import numpy as np from haystack.utils.auth import Secret from PIL.Image import Image -from sentence_transformers import SentenceTransformer +from sentence_transformers import SentenceTransformer, quantize_embeddings class _SentenceTransformersEmbeddingBackendFactory: @@ -106,4 +107,10 @@ def __init__( ) def embed(self, data: list[str] | list[Image], **kwargs: Any) -> list[list[float]]: + quantization_ranges = kwargs.pop("quantization_ranges", None) + precision = kwargs.get("precision", "float32") + if quantization_ranges is not None and precision in ("int8", "uint8"): + kwargs["precision"] = "float32" + embeddings = self.model.encode(data, **kwargs) + return quantize_embeddings(embeddings, precision=precision, ranges=np.asarray(quantization_ranges)).tolist() return self.model.encode(data, **kwargs).tolist() diff --git a/integrations/sentence_transformers/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_document_embedder.py b/integrations/sentence_transformers/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_document_embedder.py index 6af4680af7..83fdc86bfe 100644 --- a/integrations/sentence_transformers/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_document_embedder.py +++ b/integrations/sentence_transformers/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_document_embedder.py @@ -5,7 +5,7 @@ from dataclasses import replace from typing import Any, Literal -from haystack import Document, component, default_from_dict, default_to_dict +from haystack import Document, component, default_from_dict, default_to_dict, logging from haystack.utils import ComponentDevice, Secret from haystack.utils.hf import deserialize_hf_model_kwargs, serialize_hf_model_kwargs @@ -14,6 +14,8 @@ _SentenceTransformersEmbeddingBackendFactory, ) +logger = logging.getLogger(__name__) + @component class SentenceTransformersDocumentEmbedder: @@ -63,6 +65,7 @@ def __init__( encode_kwargs: dict[str, Any] | None = None, backend: Literal["torch", "onnx", "openvino"] = "torch", revision: str | None = None, + quantization_ranges: list[list[float]] | None = None, ) -> None: """ Creates a SentenceTransformersDocumentEmbedder component. @@ -124,6 +127,13 @@ def __init__( :param revision: The specific model version to use. It can be a branch name, a tag name, or a commit id, for a stored model on Hugging Face. + :param quantization_ranges: + Calibration ranges to use when `precision` is "int8" or "uint8", with shape `(2, embedding_dim)`: + minimum values in the first row and maximum values in the second. + Scalar quantization calibrates the min/max range from the batch being encoded, which is degenerate + for small batches and inconsistent across batches. Pass ranges computed from a representative + sample of embeddings to get consistent quantized embeddings, compatible with query embeddings + quantized with the same ranges. """ self.model = model @@ -147,6 +157,14 @@ def __init__( self.embedding_backend: _SentenceTransformersEmbeddingBackend | None = None self.precision = precision self.backend = backend + self.quantization_ranges = quantization_ranges + if precision in ("int8", "uint8") and quantization_ranges is None: + logger.warning( + "Using precision '{precision}' without `quantization_ranges` calibrates the quantization range " + "from each batch itself, which is degenerate for small batches and inconsistent across batches. " + "Pass `quantization_ranges` computed from a representative sample of embeddings.", + precision=precision, + ) def _get_telemetry_data(self) -> dict[str, Any]: """ @@ -183,6 +201,7 @@ def to_dict(self) -> dict[str, Any]: precision=self.precision, encode_kwargs=self.encode_kwargs, backend=self.backend, + quantization_ranges=self.quantization_ranges, ) if serialization_dict["init_parameters"].get("model_kwargs") is not None: serialize_hf_model_kwargs(serialization_dict["init_parameters"]["model_kwargs"]) @@ -262,6 +281,7 @@ def run(self, documents: list[Document]) -> dict[str, list[Document]]: show_progress_bar=self.progress_bar, normalize_embeddings=self.normalize_embeddings, precision=self.precision, + quantization_ranges=self.quantization_ranges, **(self.encode_kwargs if self.encode_kwargs else {}), ) diff --git a/integrations/sentence_transformers/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_text_embedder.py b/integrations/sentence_transformers/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_text_embedder.py index efff5af5e1..35b6dae681 100644 --- a/integrations/sentence_transformers/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_text_embedder.py +++ b/integrations/sentence_transformers/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_text_embedder.py @@ -4,7 +4,7 @@ from typing import Any, Literal -from haystack import component, default_from_dict, default_to_dict +from haystack import component, default_from_dict, default_to_dict, logging from haystack.utils import ComponentDevice, Secret from haystack.utils.hf import deserialize_hf_model_kwargs, serialize_hf_model_kwargs @@ -13,6 +13,8 @@ _SentenceTransformersEmbeddingBackendFactory, ) +logger = logging.getLogger(__name__) + @component class SentenceTransformersTextEmbedder: @@ -57,6 +59,7 @@ def __init__( encode_kwargs: dict[str, Any] | None = None, backend: Literal["torch", "onnx", "openvino"] = "torch", revision: str | None = None, + quantization_ranges: list[list[float]] | None = None, ) -> None: """ Create a SentenceTransformersTextEmbedder component. @@ -114,6 +117,12 @@ def __init__( :param revision: The specific model version to use. It can be a branch name, a tag name, or a commit id, for a stored model on Hugging Face. + :param quantization_ranges: + Calibration ranges to use when `precision` is "int8" or "uint8", with shape `(2, embedding_dim)`: + minimum values in the first row and maximum values in the second. + Scalar quantization calibrates the min/max range from the batch being encoded, which is degenerate + for a single text and produces meaningless embeddings. Pass ranges computed from a representative + sample of embeddings to get consistent quantized embeddings. """ self.model = model @@ -135,6 +144,14 @@ def __init__( self.embedding_backend: _SentenceTransformersEmbeddingBackend | None = None self.precision = precision self.backend = backend + self.quantization_ranges = quantization_ranges + if precision in ("int8", "uint8") and quantization_ranges is None: + logger.warning( + "Using precision '{precision}' without `quantization_ranges` produces meaningless embeddings for " + "single texts, because the calibration range is computed from the batch itself. " + "Pass `quantization_ranges` computed from a representative sample of embeddings.", + precision=precision, + ) def _get_telemetry_data(self) -> dict[str, Any]: """ @@ -169,6 +186,7 @@ def to_dict(self) -> dict[str, Any]: precision=self.precision, encode_kwargs=self.encode_kwargs, backend=self.backend, + quantization_ranges=self.quantization_ranges, ) if serialization_dict["init_parameters"].get("model_kwargs") is not None: serialize_hf_model_kwargs(serialization_dict["init_parameters"]["model_kwargs"]) @@ -240,6 +258,7 @@ def run(self, text: str) -> dict[str, Any]: show_progress_bar=self.progress_bar, normalize_embeddings=self.normalize_embeddings, precision=self.precision, + quantization_ranges=self.quantization_ranges, **(self.encode_kwargs if self.encode_kwargs else {}), )[0] return {"embedding": embedding} diff --git a/integrations/sentence_transformers/tests/test_sentence_transformers_document_embedder.py b/integrations/sentence_transformers/tests/test_sentence_transformers_document_embedder.py index 4e01ddf522..6366ce2f25 100644 --- a/integrations/sentence_transformers/tests/test_sentence_transformers_document_embedder.py +++ b/integrations/sentence_transformers/tests/test_sentence_transformers_document_embedder.py @@ -33,6 +33,7 @@ def test_init_default(self): assert embedder.local_files_only is False assert embedder.truncate_dim is None assert embedder.precision == "float32" + assert embedder.quantization_ranges is None def test_init_with_parameters(self): embedder = SentenceTransformersDocumentEmbedder( @@ -51,6 +52,7 @@ def test_init_with_parameters(self): local_files_only=True, truncate_dim=256, precision="int8", + quantization_ranges=[[-1.0, -1.0], [1.0, 1.0]], ) assert embedder.model == "model" assert embedder.device == ComponentDevice.from_str("cuda:0") @@ -67,6 +69,12 @@ def test_init_with_parameters(self): assert embedder.local_files_only assert embedder.truncate_dim == 256 assert embedder.precision == "int8" + assert embedder.quantization_ranges == [[-1.0, -1.0], [1.0, 1.0]] + + def test_init_quantized_precision_without_ranges_warns(self, caplog): + with caplog.at_level("WARNING"): + SentenceTransformersDocumentEmbedder(model="model", precision="int8") + assert "quantization_ranges" in caplog.text def test_to_dict(self): component = SentenceTransformersDocumentEmbedder(model="model", device=ComponentDevice.from_str("cpu")) @@ -94,6 +102,7 @@ def test_to_dict(self): "config_kwargs": None, "precision": "float32", "backend": "torch", + "quantization_ranges": None, }, } @@ -143,6 +152,7 @@ def test_to_dict_with_custom_init_parameters(self): "precision": "int8", "encode_kwargs": {"task": "clustering"}, "backend": "torch", + "quantization_ranges": None, }, } @@ -346,6 +356,7 @@ def test_embed_metadata(self): show_progress_bar=True, normalize_embeddings=False, precision="float32", + quantization_ranges=None, ) def test_embed_metadata_preserves_falsy_values(self): @@ -366,6 +377,7 @@ def test_embed_metadata_preserves_falsy_values(self): show_progress_bar=True, normalize_embeddings=False, precision="float32", + quantization_ranges=None, ) def test_embed_encode_kwargs(self): @@ -380,6 +392,7 @@ def test_embed_encode_kwargs(self): show_progress_bar=True, normalize_embeddings=False, precision="float32", + quantization_ranges=None, task="retrieval.passage", ) @@ -407,6 +420,7 @@ def test_prefix_suffix(self): show_progress_bar=True, normalize_embeddings=False, precision="float32", + quantization_ranges=None, ) @patch( diff --git a/integrations/sentence_transformers/tests/test_sentence_transformers_embedding_backend.py b/integrations/sentence_transformers/tests/test_sentence_transformers_embedding_backend.py index 6b57733cff..05b95f9e42 100644 --- a/integrations/sentence_transformers/tests/test_sentence_transformers_embedding_backend.py +++ b/integrations/sentence_transformers/tests/test_sentence_transformers_embedding_backend.py @@ -4,6 +4,7 @@ from unittest.mock import patch +import numpy as np from haystack.utils.auth import Secret from haystack_integrations.components.embedders.sentence_transformers.embedding_backend.backend import ( @@ -65,3 +66,35 @@ def test_embedding_function_with_kwargs(mock_sentence_transformer): embedding_backend.embed(data=data, normalize_embeddings=True) embedding_backend.model.encode.assert_called_once_with(data, normalize_embeddings=True) + + +@patch("haystack_integrations.components.embedders.sentence_transformers.embedding_backend.backend.quantize_embeddings") +@patch("haystack_integrations.components.embedders.sentence_transformers.embedding_backend.backend.SentenceTransformer") +def test_embedding_function_with_quantization_ranges(mock_sentence_transformer, mock_quantize_embeddings): + embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(model="quantized_model") + embedding_backend.model.encode.return_value = np.array([[0.1, 0.2]]) + mock_quantize_embeddings.return_value = np.array([[12, 34]], dtype=np.int8) + + data = ["sentence"] + ranges = [[-1.0, -1.0], [1.0, 1.0]] + result = embedding_backend.embed(data=data, precision="int8", quantization_ranges=ranges) + + embedding_backend.model.encode.assert_called_once_with(data, precision="float32") + assert mock_quantize_embeddings.call_count == 1 + _, called_kwargs = mock_quantize_embeddings.call_args + assert called_kwargs["precision"] == "int8" + assert np.array_equal(called_kwargs["ranges"], np.asarray(ranges)) + assert result == [[12, 34]] + + +@patch("haystack_integrations.components.embedders.sentence_transformers.embedding_backend.backend.quantize_embeddings") +@patch("haystack_integrations.components.embedders.sentence_transformers.embedding_backend.backend.SentenceTransformer") +def test_embedding_function_without_quantization_ranges(mock_sentence_transformer, mock_quantize_embeddings): + embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(model="another_quantized") + embedding_backend.model.encode.return_value = np.array([[1, 2]], dtype=np.int8) + + data = ["sentence"] + embedding_backend.embed(data=data, precision="int8", quantization_ranges=None) + + embedding_backend.model.encode.assert_called_once_with(data, precision="int8") + mock_quantize_embeddings.assert_not_called() diff --git a/integrations/sentence_transformers/tests/test_sentence_transformers_text_embedder.py b/integrations/sentence_transformers/tests/test_sentence_transformers_text_embedder.py index b9eeed8f1a..942259e6b5 100644 --- a/integrations/sentence_transformers/tests/test_sentence_transformers_text_embedder.py +++ b/integrations/sentence_transformers/tests/test_sentence_transformers_text_embedder.py @@ -30,6 +30,7 @@ def test_init_default(self): assert embedder.local_files_only is False assert embedder.truncate_dim is None assert embedder.precision == "float32" + assert embedder.quantization_ranges is None def test_init_with_parameters(self): embedder = SentenceTransformersTextEmbedder( @@ -46,6 +47,7 @@ def test_init_with_parameters(self): local_files_only=True, truncate_dim=256, precision="int8", + quantization_ranges=[[-1.0, -1.0], [1.0, 1.0]], ) assert embedder.model == "model" assert embedder.device == ComponentDevice.from_str("cuda:0") @@ -60,6 +62,12 @@ def test_init_with_parameters(self): assert embedder.local_files_only is True assert embedder.truncate_dim == 256 assert embedder.precision == "int8" + assert embedder.quantization_ranges == [[-1.0, -1.0], [1.0, 1.0]] + + def test_init_quantized_precision_without_ranges_warns(self, caplog): + with caplog.at_level("WARNING"): + SentenceTransformersTextEmbedder(model="model", precision="int8") + assert "quantization_ranges" in caplog.text def test_to_dict(self): component = SentenceTransformersTextEmbedder(model="model", device=ComponentDevice.from_str("cpu")) @@ -85,6 +93,7 @@ def test_to_dict(self): "config_kwargs": None, "precision": "float32", "backend": "torch", + "quantization_ranges": None, }, } @@ -129,6 +138,7 @@ def test_to_dict_with_custom_init_parameters(self): "precision": "int8", "encode_kwargs": {"task": "clustering"}, "backend": "torch", + "quantization_ranges": None, }, } @@ -300,9 +310,25 @@ def test_embed_encode_kwargs(self): show_progress_bar=True, normalize_embeddings=False, precision="float32", + quantization_ranges=None, task="retrieval.query", ) + def test_run_with_quantization_ranges(self): + ranges = [[-1.0, -1.0], [1.0, 1.0]] + embedder = SentenceTransformersTextEmbedder(model="model", precision="int8", quantization_ranges=ranges) + embedder.embedding_backend = MagicMock() + text = "a nice text to embed" + embedder.run(text=text) + embedder.embedding_backend.embed.assert_called_once_with( + [text], + batch_size=32, + show_progress_bar=True, + normalize_embeddings=False, + precision="int8", + quantization_ranges=ranges, + ) + @patch( "haystack_integrations.components.embedders.sentence_transformers.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory" ) @@ -419,3 +445,22 @@ def test_run_quantization(self, del_hf_env_vars_if_empty): assert len(embedding_def) == 128 assert all(isinstance(el, int) for el in embedding_def) + + @pytest.mark.integration + def test_run_quantization_with_ranges(self, del_hf_env_vars_if_empty): + """ + sentence-transformers-testing/stsb-bert-tiny-safetensors maps sentences & paragraphs to a 128 dimensional dense + vector space + """ + checkpoint = "sentence-transformers-testing/stsb-bert-tiny-safetensors" + text = "a nice text to embed" + ranges = [[-1.0] * 128, [1.0] * 128] + + embedder = SentenceTransformersTextEmbedder(model=checkpoint, precision="int8", quantization_ranges=ranges) + result = embedder.run(text=text) + embedding = result["embedding"] + + assert len(embedding) == 128 + assert all(isinstance(el, int) for el in embedding) + # without explicit ranges, a single text produces a degenerate all-equal embedding + assert len(set(embedding)) > 1