fix: support quantization ranges for int8 single-text embeddings

haystack/components/embedders/backends/sentence_transformers_backend.py

@@ -5,11 +5,14 @@
 import json
 from typing import Any, Literal
 
+import numpy as np
+
 from haystack.lazy_imports import LazyImport
 from haystack.utils.auth import Secret
 
 with LazyImport(message="Run 'pip install \"sentence-transformers>=5.0.0\"'") as sentence_transformers_import:
     from sentence_transformers import SentenceTransformer
+    from sentence_transformers.util import quantize_embeddings
 
 with LazyImport(message="Run 'pip install \"pillow\"'") as pillow_import:
     from PIL.Image import Image
@@ -111,4 +114,12 @@ 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"):
+            # scalar quantization calibrates min/max ranges from the batch itself, which is degenerate for
+            # small batches (a single text produces meaningless embeddings), so we quantize with explicit ranges
+            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()

haystack/components/embedders/sentence_transformers_text_embedder.py

@@ -5,14 +5,16 @@
 import warnings
 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.components.embedders.backends.sentence_transformers_backend import (
     _SentenceTransformersEmbeddingBackend,
     _SentenceTransformersEmbeddingBackendFactory,
 )
 from haystack.utils import ComponentDevice, Secret
 from haystack.utils.hf import deserialize_hf_model_kwargs, serialize_hf_model_kwargs
 
+logger = logging.getLogger(__name__)
+
 
 @component
 class SentenceTransformersTextEmbedder:
@@ -56,6 +58,7 @@ def __init__(  # noqa: PLR0913
         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.
@@ -113,6 +116,12 @@ def __init__(  # noqa: PLR0913
         :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 for "int8" and "uint8" precision, as a matrix of shape
+            (2, embedding_dim) with the minimum values in the first row and the maximum values in the second.
+            Without ranges, scalar quantization calibrates the value range from the batch itself, which is
+            degenerate for a single text and produces meaningless embeddings. Compute the ranges from a
+            representative sample of embeddings, for example the document embeddings stored in your Document Store.
         """
         warnings.warn(
             "`SentenceTransformersTextEmbedder` will be removed from Haystack in version 3.0, as it is moving to "
@@ -143,6 +152,15 @@ def __init__(  # noqa: PLR0913
         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]:
         """
@@ -177,6 +195,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"])
@@ -247,6 +266,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}

releasenotes/notes/fix-int8-single-text-zero-embeddings-bbd4074c7ce977ee.yaml

@@ -0,0 +1,10 @@
+---
+fixes:
+  - |
+    ``SentenceTransformersTextEmbedder`` produced meaningless embeddings (all zeros or all equal values)
+    when using ``precision="int8"`` or ``precision="uint8"``, because scalar quantization calibrated the
+    min/max range from the single-text batch itself. The embedder now accepts a ``quantization_ranges``
+    init parameter with explicit calibration ranges (shape ``(2, embedding_dim)``: minimum values in the
+    first row, maximum values in the second) that are forwarded to
+    ``sentence_transformers.util.quantize_embeddings``. A warning is logged when a quantized precision is
+    used without calibration ranges.

test/components/embedders/test_sentence_transformers_embedding_backend.py

@@ -4,6 +4,8 @@
 
 from unittest.mock import patch
 
+import numpy as np
+
 from haystack.components.embedders.backends.sentence_transformers_backend import (
     _SentenceTransformersEmbeddingBackendFactory,
 )
@@ -64,3 +66,34 @@ 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.components.embedders.backends.sentence_transformers_backend.quantize_embeddings")
+@patch("haystack.components.embedders.backends.sentence_transformers_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.components.embedders.backends.sentence_transformers_backend.quantize_embeddings")
+@patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer")
+def test_embedding_function_without_quantization_ranges(mock_sentence_transformer, mock_quantize_embeddings):
+    embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(model="quantized_model_2")
+
+    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()

test/components/embedders/test_sentence_transformers_text_embedder.py

@@ -28,6 +28,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(
@@ -44,6 +45,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")
@@ -58,6 +60,7 @@ 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_to_dict(self):
         component = SentenceTransformersTextEmbedder(model="model", device=ComponentDevice.from_str("cpu"))
@@ -83,6 +86,7 @@ def test_to_dict(self):
                 "config_kwargs": None,
                 "precision": "float32",
                 "backend": "torch",
+                "quantization_ranges": None,
             },
         }
 
@@ -104,6 +108,7 @@ def test_to_dict_with_custom_init_parameters(self):
             config_kwargs={"use_memory_efficient_attention": False},
             precision="int8",
             encode_kwargs={"task": "clustering"},
+            quantization_ranges=[[-1.0, -1.0], [1.0, 1.0]],
         )
         data = component.to_dict()
         assert data == {
@@ -127,6 +132,7 @@ def test_to_dict_with_custom_init_parameters(self):
                 "precision": "int8",
                 "encode_kwargs": {"task": "clustering"},
                 "backend": "torch",
+                "quantization_ranges": [[-1.0, -1.0], [1.0, 1.0]],
             },
         }
 
@@ -298,9 +304,30 @@ 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,
+        )
+
+    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
+
     @patch(
         "haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
     )
@@ -419,3 +446,22 @@ def test_run_quantization(self, del_hf_env_vars):
 
         assert len(embedding_def) == 128
         assert all(isinstance(el, int) for el in embedding_def)
+
+    @pytest.mark.integration
+    @pytest.mark.slow
+    def test_run_quantization_with_ranges(self, del_hf_env_vars):
+        """
+        Without explicit ranges, int8 quantization of a single text calibrates min/max from the batch itself,
+        producing a degenerate embedding. With ranges, the embedding must contain distinct values.
+        """
+        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)
+        assert len(set(embedding)) > 1