fix: support quantization ranges for int8/uint8 sentence-transformers…

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

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 {}),
         )
 

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}

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(

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

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