sentence-transformers-haystack 2.0.0__tar.gz → 2.1.0__tar.gz

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## [integrations/sentence_transformers-v2.0.0] - 2026-10-23
+
+### 🚜 Refactor
+
+- [**breaking**] SentenceTransformersTextEmbedder - always require quantization_ranges for int8/uint8 precision (#4154)
+
+
 ## [integrations/sentence_transformers-v1.0.0] - 2026-10-16
 
 ### 🧹 Chores

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: sentence-transformers-haystack
-Version: 2.0.0
+Version: 2.1.0
 Summary: Haystack integration for sentence_transformers
 Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/sentence_transformers#readme
 Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/pyproject.toml

@@ -57,6 +57,7 @@ dependencies = [
     "pytest-rerunfailures",
     "mypy",
     "pip",
+    "onnxruntime",
 ]
 
 [tool.hatch.envs.test.scripts]
@@ -78,6 +79,7 @@ disallow_incomplete_defs = true
 module = [
   # these libraries do not ship type stubs
   "PIL.*",
+  "onnxruntime",
 ]
 ignore_missing_imports = true
 

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/src/haystack_integrations/components/embedders/sentence_transformers/embedding_backend/backend.py

@@ -6,11 +6,36 @@ import json
 from typing import Any, ClassVar, Literal
 
 import numpy as np
+from haystack.lazy_imports import LazyImport
 from haystack.utils.auth import Secret
 from PIL.Image import Image
 
 from sentence_transformers import SentenceTransformer, quantize_embeddings
 
+with LazyImport(message="Run 'pip install \"sentence-transformers[onnx]\"'") as onnxruntime_import:
+    import onnxruntime as ort
+
+
+def _resolve_onnx_session_options(model_kwargs: dict[str, Any] | None) -> dict[str, Any] | None:
+    """
+    Build an `onnxruntime.SessionOptions` object from a `session_options` dict in `model_kwargs`.
+
+    Each key of the dict is set as an attribute of the `SessionOptions`, e.g. `{"intra_op_num_threads": 1}`.
+
+    :param model_kwargs: The model kwargs passed to the Sentence Transformers model.
+    :returns: A copy of `model_kwargs` with `session_options` replaced by the built `SessionOptions` object, or
+        `model_kwargs` itself when it has no `session_options` dict.
+    :raises ImportError: If `model_kwargs` has a `session_options` dict and `onnxruntime` is not installed.
+    """
+    if not model_kwargs or not isinstance(model_kwargs.get("session_options"), dict):
+        return model_kwargs
+
+    onnxruntime_import.check()
+    session_options = ort.SessionOptions()
+    for name, value in model_kwargs["session_options"].items():
+        setattr(session_options, name, value)
+    return {**model_kwargs, "session_options": session_options}
+
 
 class _SentenceTransformersEmbeddingBackendFactory:
     """
@@ -99,7 +124,9 @@ class _SentenceTransformersEmbeddingBackend:
             revision=revision,
             local_files_only=local_files_only,
             truncate_dim=truncate_dim,
-            model_kwargs=model_kwargs,
+            model_kwargs=_resolve_onnx_session_options(model_kwargs=model_kwargs)
+            if backend == "onnx"
+            else model_kwargs,
             # `tokenizer_kwargs` was renamed to `processor_kwargs` in sentence-transformers 5.4.0
             processor_kwargs=tokenizer_kwargs,
             config_kwargs=config_kwargs,

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/src/haystack_integrations/components/embedders/sentence_transformers/embedding_backend/sparse_backend.py

@@ -8,6 +8,9 @@ from typing import Any, ClassVar, Literal
 from haystack.dataclasses.sparse_embedding import SparseEmbedding
 from haystack.utils.auth import Secret
 
+from haystack_integrations.components.embedders.sentence_transformers.embedding_backend.backend import (
+    _resolve_onnx_session_options,
+)
 from sentence_transformers import SparseEncoder
 
 
@@ -93,7 +96,9 @@ class _SentenceTransformersSparseEncoderEmbeddingBackend:
             trust_remote_code=trust_remote_code,
             revision=revision,
             local_files_only=local_files_only,
-            model_kwargs=model_kwargs,
+            model_kwargs=_resolve_onnx_session_options(model_kwargs=model_kwargs)
+            if backend == "onnx"
+            else model_kwargs,
             # `tokenizer_kwargs` was renamed to `processor_kwargs` in sentence-transformers 5.4.0
             processor_kwargs=tokenizer_kwargs,
             config_kwargs=config_kwargs,

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_doc_image_embedder.py

@@ -113,6 +113,9 @@ class SentenceTransformersDocumentImageEmbedder:
         :param model_kwargs:
             Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
             when loading the model. Refer to specific model documentation for available kwargs.
+            With `backend="onnx"`, a `session_options` dict is turned into an `onnxruntime.SessionOptions` with
+            those attributes set, for example `{"session_options": {"intra_op_num_threads": 1}}` to limit
+            the threads each inference call uses.
         :param tokenizer_kwargs:
             Additional keyword arguments forwarded as `processor_kwargs` to the Sentence Transformers model
             when loading it. Refer to the

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_document_embedder.py

@@ -106,6 +106,9 @@ class SentenceTransformersDocumentEmbedder:
         :param model_kwargs:
             Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
             when loading the model. Refer to specific model documentation for available kwargs.
+            With `backend="onnx"`, a `session_options` dict is turned into an `onnxruntime.SessionOptions` with
+            those attributes set, for example `{"session_options": {"intra_op_num_threads": 1}}` to limit
+            the threads each inference call uses.
         :param tokenizer_kwargs:
             Additional keyword arguments forwarded as `processor_kwargs` to the Sentence Transformers model
             when loading it. Refer to the

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_sparse_document_embedder.py

@@ -94,6 +94,9 @@ class SentenceTransformersSparseDocumentEmbedder:
         :param model_kwargs:
             Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
             when loading the model. Refer to specific model documentation for available kwargs.
+            With `backend="onnx"`, a `session_options` dict is turned into an `onnxruntime.SessionOptions` with
+            those attributes set, for example `{"session_options": {"intra_op_num_threads": 1}}` to limit
+            the threads each inference call uses.
         :param tokenizer_kwargs:
             Additional keyword arguments forwarded as `processor_kwargs` to the Sentence Transformers model
             when loading it. Refer to the

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_sparse_text_embedder.py

@@ -75,6 +75,9 @@ class SentenceTransformersSparseTextEmbedder:
         :param model_kwargs:
             Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
             when loading the model. Refer to specific model documentation for available kwargs.
+            With `backend="onnx"`, a `session_options` dict is turned into an `onnxruntime.SessionOptions` with
+            those attributes set, for example `{"session_options": {"intra_op_num_threads": 1}}` to limit
+            the threads each inference call uses.
         :param tokenizer_kwargs:
             Additional keyword arguments forwarded as `processor_kwargs` to the Sentence Transformers model
             when loading it. Refer to the

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/src/haystack_integrations/components/embedders/sentence_transformers/sentence_transformers_text_embedder.py

@@ -94,6 +94,9 @@ class SentenceTransformersTextEmbedder:
         :param model_kwargs:
             Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
             when loading the model. Refer to specific model documentation for available kwargs.
+            With `backend="onnx"`, a `session_options` dict is turned into an `onnxruntime.SessionOptions` with
+            those attributes set, for example `{"session_options": {"intra_op_num_threads": 1}}` to limit
+            the threads each inference call uses.
         :param tokenizer_kwargs:
             Additional keyword arguments forwarded as `processor_kwargs` to the Sentence Transformers model
             when loading it. Refer to the

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/src/haystack_integrations/components/rankers/sentence_transformers/sentence_transformers_diversity.py

@@ -11,6 +11,9 @@ from haystack.utils import ComponentDevice, Secret
 from haystack.utils.hf import deserialize_hf_model_kwargs, serialize_hf_model_kwargs
 from haystack.utils.misc import _deduplicate_documents
 
+from haystack_integrations.components.embedders.sentence_transformers.embedding_backend.backend import (
+    _resolve_onnx_session_options,
+)
 from sentence_transformers import SentenceTransformer
 
 
@@ -162,6 +165,9 @@ class SentenceTransformersDiversityRanker:
         :param model_kwargs:
             Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
             when loading the model. Refer to specific model documentation for available kwargs.
+            With `backend="onnx"`, a `session_options` dict is turned into an `onnxruntime.SessionOptions` with
+            those attributes set, for example `{"session_options": {"intra_op_num_threads": 1}}` to limit
+            the threads each inference call uses.
         :param tokenizer_kwargs:
             Additional keyword arguments forwarded as `processor_kwargs` to the Sentence Transformers model
             when loading it. Refer to the
@@ -206,7 +212,11 @@ class SentenceTransformersDiversityRanker:
                 model_name_or_path=self.model_name_or_path,
                 device=self.device.to_torch_str(),
                 token=self.token.resolve_value() if self.token else None,
-                model_kwargs=self.model_kwargs,
+                model_kwargs=(
+                    _resolve_onnx_session_options(model_kwargs=self.model_kwargs)
+                    if self.backend == "onnx"
+                    else self.model_kwargs
+                ),
                 # `tokenizer_kwargs` was renamed to `processor_kwargs` in sentence-transformers 5.4.0
                 processor_kwargs=self.tokenizer_kwargs,
                 config_kwargs=self.config_kwargs,

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/src/haystack_integrations/components/rankers/sentence_transformers/sentence_transformers_similarity.py

@@ -12,6 +12,9 @@ from haystack.utils.hf import deserialize_hf_model_kwargs, serialize_hf_model_kw
 from haystack.utils.misc import _deduplicate_documents
 from torch.nn import Identity, Sigmoid
 
+from haystack_integrations.components.embedders.sentence_transformers.embedding_backend.backend import (
+    _resolve_onnx_session_options,
+)
 from sentence_transformers import CrossEncoder
 
 
@@ -97,6 +100,9 @@ class SentenceTransformersSimilarityRanker:
         :param model_kwargs:
             Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
             when loading the model. Refer to specific model documentation for available kwargs.
+            With `backend="onnx"`, a `session_options` dict is turned into an `onnxruntime.SessionOptions` with
+            those attributes set, for example `{"session_options": {"intra_op_num_threads": 1}}` to limit
+            the threads each inference call uses.
         :param tokenizer_kwargs:
             Additional keyword arguments forwarded as `processor_kwargs` to the Sentence Transformers model
             when loading it. Refer to the
@@ -155,7 +161,11 @@ class SentenceTransformersSimilarityRanker:
                 device=self.device.to_torch_str(),
                 token=self.token.resolve_value() if self.token else None,
                 trust_remote_code=self.trust_remote_code,
-                model_kwargs=self.model_kwargs,
+                model_kwargs=(
+                    _resolve_onnx_session_options(model_kwargs=self.model_kwargs)
+                    if self.backend == "onnx"
+                    else self.model_kwargs
+                ),
                 # `tokenizer_kwargs` was renamed to `processor_kwargs` in sentence-transformers 5.4.0
                 processor_kwargs=self.tokenizer_kwargs,
                 config_kwargs=self.config_kwargs,

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/tests/test_sentence_transformers_diversity.py

@@ -4,6 +4,7 @@
 
 from unittest.mock import MagicMock, call, patch
 
+import onnxruntime as ort
 import pytest
 import torch
 from haystack import Document, Pipeline
@@ -741,6 +742,26 @@ class TestSentenceTransformersDiversityRanker:
             backend="onnx",
         )
 
+    @patch(
+        "haystack_integrations.components.rankers.sentence_transformers.sentence_transformers_diversity.SentenceTransformer"
+    )
+    def test_from_dict_warm_up_onnx_session_options(self, mocked_sentence_transformer):
+        ranker = SentenceTransformersDiversityRanker.from_dict(
+            {
+                "type": "haystack_integrations.components.rankers.sentence_transformers.sentence_transformers_diversity.SentenceTransformersDiversityRanker",
+                "init_parameters": {
+                    "backend": "onnx",
+                    "model_kwargs": {"session_options": {"intra_op_num_threads": 1}},
+                },
+            }
+        )
+        ranker.warm_up()
+
+        session_options = mocked_sentence_transformer.call_args.kwargs["model_kwargs"]["session_options"]
+        assert isinstance(session_options, ort.SessionOptions)
+        assert session_options.intra_op_num_threads == 1
+        assert ranker.to_dict()["init_parameters"]["model_kwargs"]["session_options"] == {"intra_op_num_threads": 1}
+
     @patch(
         "haystack_integrations.components.rankers.sentence_transformers.sentence_transformers_diversity.SentenceTransformer"
     )

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/tests/test_sentence_transformers_embedding_backend.py

@@ -5,9 +5,11 @@
 from unittest.mock import patch
 
 import numpy as np
+import onnxruntime as ort
 from haystack.utils.auth import Secret
 
 from haystack_integrations.components.embedders.sentence_transformers.embedding_backend.backend import (
+    _resolve_onnx_session_options,
     _SentenceTransformersEmbeddingBackendFactory,
 )
 
@@ -58,6 +60,29 @@ def test_model_initialization(mock_sentence_transformer):
     )
 
 
+def test_resolve_onnx_session_options():
+    model_kwargs = {"file_name": "onnx/model.onnx", "session_options": {"intra_op_num_threads": 1}}
+
+    resolved = _resolve_onnx_session_options(model_kwargs=model_kwargs)
+
+    assert isinstance(resolved["session_options"], ort.SessionOptions)
+    assert resolved["session_options"].intra_op_num_threads == 1
+    assert resolved["file_name"] == "onnx/model.onnx"
+    # the component's own model_kwargs stay serializable
+    assert model_kwargs["session_options"] == {"intra_op_num_threads": 1}
+
+
+@patch("haystack_integrations.components.embedders.sentence_transformers.embedding_backend.backend.SentenceTransformer")
+def test_model_initialization_onnx_session_options(mock_sentence_transformer):
+    _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
+        model="onnx_session_model", model_kwargs={"session_options": {"intra_op_num_threads": 2}}, backend="onnx"
+    )
+
+    session_options = mock_sentence_transformer.call_args.kwargs["model_kwargs"]["session_options"]
+    assert isinstance(session_options, ort.SessionOptions)
+    assert session_options.intra_op_num_threads == 2
+
+
 @patch("haystack_integrations.components.embedders.sentence_transformers.embedding_backend.backend.SentenceTransformer")
 def test_embedding_function_with_kwargs(mock_sentence_transformer):
     embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(model="model")

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/tests/test_sentence_transformers_similarity.py

@@ -5,6 +5,7 @@
 from unittest.mock import MagicMock, patch
 
 import numpy as np
+import onnxruntime as ort
 import pytest
 import torch
 from haystack import Document
@@ -68,6 +69,26 @@ class TestSentenceTransformersSimilarityRanker:
             backend="onnx",
         )
 
+    @patch(
+        "haystack_integrations.components.rankers.sentence_transformers.sentence_transformers_similarity.CrossEncoder"
+    )
+    def test_from_dict_warm_up_onnx_session_options(self, mocked_cross_encoder):
+        ranker = SentenceTransformersSimilarityRanker.from_dict(
+            {
+                "type": "haystack_integrations.components.rankers.sentence_transformers.sentence_transformers_similarity.SentenceTransformersSimilarityRanker",
+                "init_parameters": {
+                    "backend": "onnx",
+                    "model_kwargs": {"session_options": {"intra_op_num_threads": 1}},
+                },
+            }
+        )
+        ranker.warm_up()
+
+        session_options = mocked_cross_encoder.call_args.kwargs["model_kwargs"]["session_options"]
+        assert isinstance(session_options, ort.SessionOptions)
+        assert session_options.intra_op_num_threads == 1
+        assert ranker.to_dict()["init_parameters"]["model_kwargs"]["session_options"] == {"intra_op_num_threads": 1}
+
     @patch(
         "haystack_integrations.components.rankers.sentence_transformers.sentence_transformers_similarity.CrossEncoder"
     )

{sentence_transformers_haystack-2.0.0 → sentence_transformers_haystack-2.1.0}/tests/test_sentence_transformers_sparse_embedding_backend.py

@@ -4,6 +4,7 @@
 
 from unittest.mock import patch
 
+import onnxruntime as ort
 import pytest
 import torch
 from haystack.dataclasses.sparse_embedding import SparseEmbedding
@@ -62,6 +63,19 @@ def test_sparse_model_initialization(mock_sparse_encoder):
     )
 
 
+@patch(
+    "haystack_integrations.components.embedders.sentence_transformers.embedding_backend.sparse_backend.SparseEncoder"
+)
+def test_sparse_model_initialization_onnx_session_options(mock_sparse_encoder):
+    _SentenceTransformersSparseEmbeddingBackendFactory.get_embedding_backend(
+        model="onnx_session_model", model_kwargs={"session_options": {"intra_op_num_threads": 2}}, backend="onnx"
+    )
+
+    session_options = mock_sparse_encoder.call_args.kwargs["model_kwargs"]["session_options"]
+    assert isinstance(session_options, ort.SessionOptions)
+    assert session_options.intra_op_num_threads == 2
+
+
 @patch(
     "haystack_integrations.components.embedders.sentence_transformers.embedding_backend.sparse_backend.SparseEncoder"
 )