deepset-ai/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2026-10-12 02:12:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore!: remove sentence transformers components (#11646)

Browse Source
This commit is contained in:
Stefano Fiorucci
2026-06-16 12:23:49 +02:00
committed by GitHub
parent 496c92a73b
commit 6411c8cef9
48 changed files with 185 additions and 6257 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/slow.yml
-14
View File
@@ -59,30 +59,16 @@ jobs:
changes:
- "haystack/components/audio/whisper_local.py"
- "haystack/components/converters/tika.py"
- "haystack/components/embedders/backends/sentence_transformers_backend.py"
- "haystack/components/embedders/backends/sentence_transformers_sparse_backend.py"
- "haystack/components/embedders/image/sentence_transformers_doc_image_embedder.py"
- "haystack/components/embedders/sentence_transformers_text_embedder.py"
- "haystack/components/embedders/sentence_transformers_sparse_document_embedder.py"
- "haystack/components/embedders/sentence_transformers_sparse_text_embedder.py"
- "haystack/components/evaluators/sas_evaluator.py"
- "haystack/components/generators/openai_dalle.py"
- "haystack/components/preprocessors/embedding_based_document_splitter.py"
- "haystack/components/rankers/sentence_transformers_diversity.py"
- "haystack/components/rankers/sentence_transformers_similarity.py"
- "haystack/components/retrievers/multi_query_embedding_retriever.py"
- "test/components/audio/test_whisper_local.py"
- "test/components/converters/test_tika_doc_converter.py"
- "test/components/embedders/image/test_sentence_transformers_doc_image_embedder.py"
- "test/components/embedders/test_sentence_transformers_text_embedder.py"
- "test/components/embedders/test_sentence_transformers_sparse_document_embedder.py"
- "test/components/embedders/test_sentence_transformers_sparse_text_embedder.py"
- "test/components/evaluators/test_sas_evaluator.py"
- "test/components/generators/test_openai_dalle.py"
- "test/components/preprocessors/test_embedding_based_document_splitter.py"
- "test/components/rankers/test_sentence_transformers_diversity.py"
- "test/components/rankers/test_sentence_transformers_similarity.py"
- "test/components/retrievers/test_multi_query_embedding_retriever.py"
slow-integration-tests:
MIGRATION.md
+7
View File
@@ -101,6 +101,13 @@ pip install <new-package>
| `from haystack.components.websearch import SerperDevWebSearch` | `serperdev-haystack` | `from haystack_integrations.components.websearch.serperdev import SerperDevWebSearch` |
| `from haystack.components.websearch import SearchApiWebSearch` | `searchapi-haystack` | `from haystack_integrations.components.websearch.searchapi import SearchApiWebSearch` || `from haystack.components.extractors import NamedEntityExtractor` (Hugging Face backend) | `transformers-haystack` | `from haystack_integrations.components.extractors.transformers import TransformersNamedEntityExtractor` |
| `from haystack.components.extractors import NamedEntityExtractor` (spaCy backend) | `spacy-haystack` | `from haystack_integrations.components.extractors.spacy import SpacyNamedEntityExtractor` |
| `from haystack.components.embedders import SentenceTransformersTextEmbedder` | `sentence-transformers-haystack` | `from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersTextEmbedder` |
| `from haystack.components.embedders import SentenceTransformersDocumentEmbedder` | `sentence-transformers-haystack` | `from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersDocumentEmbedder` |
| `from haystack.components.embedders import SentenceTransformersSparseTextEmbedder` | `sentence-transformers-haystack` | `from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersSparseTextEmbedder` |
| `from haystack.components.embedders import SentenceTransformersSparseDocumentEmbedder` | `sentence-transformers-haystack` | `from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersSparseDocumentEmbedder` |
| `from haystack.components.embedders.image import SentenceTransformersDocumentImageEmbedder` | `sentence-transformers-haystack` | `from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersDocumentImageEmbedder` |
| `from haystack.components.rankers import SentenceTransformersSimilarityRanker` | `sentence-transformers-haystack` | `from haystack_integrations.components.rankers.sentence_transformers import SentenceTransformersSimilarityRanker` |
| `from haystack.components.rankers import SentenceTransformersDiversityRanker` | `sentence-transformers-haystack` | `from haystack_integrations.components.rankers.sentence_transformers import SentenceTransformersDiversityRanker` |
### `TransformersSimilarityRanker` removed
e2e/pipelines/test_dense_doc_search.py
-83
View File
@@ -1,83 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
import json
from haystack import Pipeline
from haystack.components.converters import PyPDFToDocument, TextFileToDocument
from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
from haystack.components.joiners import DocumentJoiner
from haystack.components.preprocessors import DocumentCleaner, DocumentSplitter
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from haystack.components.routers import FileTypeRouter
from haystack.components.writers import DocumentWriter
from haystack.document_stores.in_memory import InMemoryDocumentStore
def test_dense_doc_search_pipeline(tmp_path, samples_path, del_hf_env_vars):
# Create the indexing pipeline
indexing_pipeline = Pipeline()
indexing_pipeline.add_component(
instance=FileTypeRouter(mime_types=["text/plain", "application/pdf"]), name="file_type_router"
)
indexing_pipeline.add_component(instance=TextFileToDocument(), name="text_file_converter")
indexing_pipeline.add_component(instance=PyPDFToDocument(), name="pdf_file_converter")
indexing_pipeline.add_component(instance=DocumentJoiner(), name="joiner")
indexing_pipeline.add_component(instance=DocumentCleaner(), name="cleaner")
indexing_pipeline.add_component(
instance=DocumentSplitter(split_by="period", split_length=250, split_overlap=30), name="splitter"
)
indexing_pipeline.add_component(
instance=SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"), name="embedder"
)
indexing_pipeline.add_component(instance=DocumentWriter(document_store=InMemoryDocumentStore()), name="writer")
indexing_pipeline.connect("file_type_router.text/plain", "text_file_converter.sources")
indexing_pipeline.connect("file_type_router.application/pdf", "pdf_file_converter.sources")
indexing_pipeline.connect("text_file_converter.documents", "joiner.documents")
indexing_pipeline.connect("pdf_file_converter.documents", "joiner.documents")
indexing_pipeline.connect("joiner.documents", "cleaner.documents")
indexing_pipeline.connect("cleaner.documents", "splitter.documents")
indexing_pipeline.connect("splitter.documents", "embedder.documents")
indexing_pipeline.connect("embedder.documents", "writer.documents")
# Serialize the indexing pipeline to YAML.
with open(tmp_path / "test_dense_doc_search_indexing_pipeline.yaml", "w") as f:
indexing_pipeline.dump(f)
# Load the indexing pipeline back
with open(tmp_path / "test_dense_doc_search_indexing_pipeline.yaml", "r") as f:
indexing_pipeline = Pipeline.load(f)
indexing_result = indexing_pipeline.run({"file_type_router": {"sources": list(samples_path.iterdir())}})
filled_document_store = indexing_pipeline.get_component("writer").document_store
assert indexing_result["writer"]["documents_written"] == 2
assert filled_document_store.count_documents() == 2
# Create the querying pipeline
query_pipeline = Pipeline()
query_pipeline.add_component(
instance=SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"), name="text_embedder"
)
query_pipeline.add_component(
instance=InMemoryEmbeddingRetriever(document_store=filled_document_store, top_k=20), name="embedding_retriever"
)
query_pipeline.connect("text_embedder", "embedding_retriever")
querying_result = query_pipeline.run({"text_embedder": {"text": "Who lives in Rome?"}})
assert querying_result["embedding_retriever"]["documents"][0].content == "My name is Giorgio and I live in Rome."
# Serialize the querying pipeline to JSON
with open(tmp_path / "test_dense_doc_search_query_pipeline.json", "w") as f:
print(json.dumps(query_pipeline.to_dict(), indent=4))
json.dump(query_pipeline.to_dict(), f)
# Load the querying pipeline back
with open(tmp_path / "test_dense_doc_search_query_pipeline.json", "r") as f:
query_pipeline = Pipeline.from_dict(json.load(f))
# Run the reloaded pipeline to verify the serialization round-trip preserves behavior
querying_result = query_pipeline.run({"text_embedder": {"text": "Who lives in Rome?"}})
assert querying_result["embedding_retriever"]["documents"][0].content == "My name is Giorgio and I live in Rome."
e2e/pipelines/test_evaluation_pipeline.py
+5 -5
View File
@@ -8,7 +8,7 @@ import pytest
from haystack import Document, Pipeline
from haystack.components.builders import AnswerBuilder, ChatPromptBuilder
from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
from haystack.components.embedders import OpenAIDocumentEmbedder, OpenAITextEmbedder
from haystack.components.evaluators import (
ContextRelevanceEvaluator,
DocumentMAPEvaluator,
@@ -26,14 +26,14 @@ from haystack.document_stores.types import DuplicatePolicy
from haystack.evaluation import EvaluationRunResult
from haystack.dataclasses import ChatMessage
EMBEDDINGS_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
EMBEDDINGS_MODEL = "text-embedding-3-small"
def indexing_pipeline(documents: list[Document]):
"""Indexing the documents"""
document_store = InMemoryDocumentStore()
doc_writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.SKIP)
doc_embedder = SentenceTransformersDocumentEmbedder(model=EMBEDDINGS_MODEL, progress_bar=False)
doc_embedder = OpenAIDocumentEmbedder(model=EMBEDDINGS_MODEL, progress_bar=False)
ingestion_pipe = Pipeline()
ingestion_pipe.add_component(instance=doc_embedder, name="doc_embedder")
ingestion_pipe.add_component(instance=doc_writer, name="doc_writer")
@@ -58,7 +58,7 @@ def rag_pipeline(document_store: InMemoryDocumentStore, top_k: int):
),
]
rag = Pipeline()
rag.add_component("embedder", SentenceTransformersTextEmbedder(model=EMBEDDINGS_MODEL, progress_bar=False))
rag.add_component("embedder", OpenAITextEmbedder(model=EMBEDDINGS_MODEL))
rag.add_component("retriever", InMemoryEmbeddingRetriever(document_store, top_k=top_k))
rag.add_component("prompt_builder", ChatPromptBuilder(template=template))
rag.add_component("generator", OpenAIChatGenerator(model="gpt-4o-mini"))
@@ -86,7 +86,7 @@ def evaluation_pipeline():
eval_pipeline = Pipeline()
eval_pipeline.add_component("doc_mrr", DocumentMRREvaluator())
eval_pipeline.add_component("groundedness", FaithfulnessEvaluator())
eval_pipeline.add_component("sas", SASEvaluator(model=EMBEDDINGS_MODEL))
eval_pipeline.add_component("sas", SASEvaluator())
eval_pipeline.add_component("doc_map", DocumentMAPEvaluator())
eval_pipeline.add_component("doc_recall_single_hit", DocumentRecallEvaluator(mode=RecallMode.SINGLE_HIT))
eval_pipeline.add_component("doc_recall_multi_hit", DocumentRecallEvaluator(mode=RecallMode.MULTI_HIT))
e2e/pipelines/test_hybrid_doc_search_pipeline.py
-58
View File
@@ -1,58 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from haystack import Document, Pipeline
from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
from haystack.components.joiners.document_joiner import DocumentJoiner
from haystack.components.rankers import SentenceTransformersSimilarityRanker
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever, InMemoryEmbeddingRetriever
from haystack.document_stores.in_memory import InMemoryDocumentStore
def test_hybrid_doc_search_pipeline(tmp_path, del_hf_env_vars):
# Create the pipeline
document_store = InMemoryDocumentStore()
hybrid_pipeline = Pipeline()
hybrid_pipeline.add_component(instance=InMemoryBM25Retriever(document_store=document_store), name="bm25_retriever")
hybrid_pipeline.add_component(
instance=SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"), name="text_embedder"
)
hybrid_pipeline.add_component(
instance=InMemoryEmbeddingRetriever(document_store=document_store), name="embedding_retriever"
)
hybrid_pipeline.add_component(instance=DocumentJoiner(), name="joiner")
hybrid_pipeline.add_component(instance=SentenceTransformersSimilarityRanker(top_k=20), name="ranker")
hybrid_pipeline.connect("bm25_retriever", "joiner")
hybrid_pipeline.connect("text_embedder", "embedding_retriever")
hybrid_pipeline.connect("embedding_retriever", "joiner")
hybrid_pipeline.connect("joiner", "ranker")
# Serialize the pipeline to YAML
with open(tmp_path / "test_hybrid_doc_search_pipeline.yaml", "w") as f:
hybrid_pipeline.dump(f)
# Load the pipeline back
with open(tmp_path / "test_hybrid_doc_search_pipeline.yaml", "r") as f:
hybrid_pipeline = Pipeline.load(f)
# Populate the document store
documents = [
Document(content="My name is Jean and I live in Paris."),
Document(content="My name is Mark and I live in Berlin."),
Document(content="My name is Mario and I live in the capital of Italy."),
Document(content="My name is Giorgio and I live in Rome."),
]
doc_embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
doc_embedder.warm_up()
embedded_documents = doc_embedder.run(documents=documents)["documents"]
hybrid_pipeline.get_component("embedding_retriever").document_store.write_documents(embedded_documents)
query = "Who lives in Rome?"
result = hybrid_pipeline.run(
{"bm25_retriever": {"query": query}, "text_embedder": {"text": query}, "ranker": {"query": query}}
)
assert result["ranker"]["documents"][0].content == "My name is Giorgio and I live in Rome."
assert result["ranker"]["documents"][1].content == "My name is Mario and I live in the capital of Italy."
e2e/pipelines/test_preprocessing_pipeline.py
+3 -2
View File
@@ -5,7 +5,7 @@
from haystack import Pipeline
from haystack.components.classifiers import DocumentLanguageClassifier
from haystack.components.converters import TextFileToDocument
from haystack.components.embedders import SentenceTransformersDocumentEmbedder
from haystack.components.embedders import OpenAIDocumentEmbedder
from haystack.components.preprocessors import DocumentCleaner, DocumentSplitter
from haystack.components.routers import FileTypeRouter, MetadataRouter
from haystack.components.writers import DocumentWriter
@@ -25,7 +25,7 @@ def test_preprocessing_pipeline(tmp_path, del_hf_env_vars):
preprocessing_pipeline.add_component(instance=DocumentCleaner(), name="cleaner")
preprocessing_pipeline.add_component(instance=DocumentSplitter(split_by="period", split_length=1), name="splitter")
preprocessing_pipeline.add_component(
instance=SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"), name="embedder"
instance=OpenAIDocumentEmbedder(model="text-embedding-3-small"), name="embedder"
)
preprocessing_pipeline.add_component(instance=DocumentWriter(document_store=document_store), name="writer")
preprocessing_pipeline.connect("file_type_router.text/plain", "text_file_converter.sources")
@@ -82,3 +82,4 @@ def test_preprocessing_pipeline(tmp_path, del_hf_env_vars):
]
assert expected_texts == [document.content for document in stored_documents]
assert all(document.meta["language"] == "en" for document in stored_documents)
assert all(document.embedding is not None for document in stored_documents)
e2e/pipelines/test_rag_pipelines_e2e.py
-152
View File
@@ -1,152 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
import json
import os
import pytest
from haystack import Document, Pipeline
from haystack.components.builders.answer_builder import AnswerBuilder
from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder
from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever, InMemoryEmbeddingRetriever
from haystack.components.writers import DocumentWriter
from haystack.dataclasses import ChatMessage
from haystack.document_stores.in_memory import InMemoryDocumentStore
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
def test_bm25_rag_pipeline(tmp_path):
# Create the RAG pipeline
prompt_template = [
ChatMessage.from_user(
"Given these documents, answer the question.\nDocuments:\n"
"{% for doc in documents %}\n {{ doc.content }}\n{% endfor %}\n"
"\nQuestion: {{question}}\nAnswer:"
)
]
rag_pipeline = Pipeline()
rag_pipeline.add_component(instance=InMemoryBM25Retriever(document_store=InMemoryDocumentStore()), name="retriever")
rag_pipeline.add_component(instance=ChatPromptBuilder(template=prompt_template), name="prompt_builder")
rag_pipeline.add_component(instance=OpenAIChatGenerator(), name="llm")
rag_pipeline.add_component(instance=AnswerBuilder(), name="answer_builder")
rag_pipeline.connect("retriever", "prompt_builder.documents")
rag_pipeline.connect("prompt_builder.prompt", "llm.messages")
rag_pipeline.connect("llm.replies", "answer_builder.replies")
rag_pipeline.connect("retriever", "answer_builder.documents")
# Serialize the pipeline to YAML
with open(tmp_path / "test_bm25_rag_pipeline.yaml", "w") as f:
rag_pipeline.dump(f)
# Load the pipeline back
with open(tmp_path / "test_bm25_rag_pipeline.yaml", "r") as f:
rag_pipeline = Pipeline.load(f)
# Populate the document store
documents = [
Document(content="My name is Jean and I live in Paris."),
Document(content="My name is Mark and I live in Berlin."),
Document(content="My name is Giorgio and I live in Rome."),
]
rag_pipeline.get_component("retriever").document_store.write_documents(documents)
# Query and assert
questions = ["Who lives in Paris?", "Who lives in Berlin?", "Who lives in Rome?"]
answers_spywords = ["Jean", "Mark", "Giorgio"]
for question, spyword in zip(questions, answers_spywords):
result = rag_pipeline.run(
{
"retriever": {"query": question},
"prompt_builder": {"question": question},
"answer_builder": {"query": question},
}
)
assert len(result["answer_builder"]["answers"]) == 1
generated_answer = result["answer_builder"]["answers"][0]
assert spyword in generated_answer.data
assert generated_answer.query == question
assert hasattr(generated_answer, "documents")
assert hasattr(generated_answer, "meta")
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
def test_embedding_retrieval_rag_pipeline(tmp_path):
# Create the RAG pipeline
prompt_template = [
ChatMessage.from_user(
"Given these documents, answer the question.\nDocuments:\n"
"{% for doc in documents %}\n {{ doc.content }}\n{% endfor %}\n"
"\nQuestion: {{question}}\nAnswer:"
)
]
rag_pipeline = Pipeline()
rag_pipeline.add_component(
instance=SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"), name="text_embedder"
)
rag_pipeline.add_component(
instance=InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore()), name="retriever"
)
rag_pipeline.add_component(instance=ChatPromptBuilder(template=prompt_template), name="prompt_builder")
rag_pipeline.add_component(instance=OpenAIChatGenerator(), name="llm")
rag_pipeline.add_component(instance=AnswerBuilder(), name="answer_builder")
rag_pipeline.connect("text_embedder", "retriever")
rag_pipeline.connect("retriever", "prompt_builder.documents")
rag_pipeline.connect("prompt_builder.prompt", "llm.messages")
rag_pipeline.connect("llm.replies", "answer_builder.replies")
rag_pipeline.connect("retriever", "answer_builder.documents")
# Serialize the pipeline to JSON
with open(tmp_path / "test_embedding_rag_pipeline.json", "w") as f:
json.dump(rag_pipeline.to_dict(), f)
# Load the pipeline back
with open(tmp_path / "test_embedding_rag_pipeline.json", "r") as f:
rag_pipeline = Pipeline.from_dict(json.load(f))
# Populate the document store
documents = [
Document(content="My name is Jean and I live in Paris."),
Document(content="My name is Mark and I live in Berlin."),
Document(content="My name is Giorgio and I live in Rome."),
]
document_store = rag_pipeline.get_component("retriever").document_store
indexing_pipeline = Pipeline()
indexing_pipeline.add_component(
instance=SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"),
name="document_embedder",
)
indexing_pipeline.add_component(instance=DocumentWriter(document_store=document_store), name="document_writer")
indexing_pipeline.connect("document_embedder", "document_writer")
indexing_pipeline.run({"document_embedder": {"documents": documents}})
# Query and assert
questions = ["Who lives in Paris?", "Who lives in Berlin?", "Who lives in Rome?"]
answers_spywords = ["Jean", "Mark", "Giorgio"]
for question, spyword in zip(questions, answers_spywords):
result = rag_pipeline.run(
{
"text_embedder": {"text": question},
"prompt_builder": {"question": question},
"answer_builder": {"query": question},
}
)
assert len(result["answer_builder"]["answers"]) == 1
generated_answer = result["answer_builder"]["answers"][0]
assert spyword in generated_answer.data
assert generated_answer.query == question
assert hasattr(generated_answer, "documents")
assert hasattr(generated_answer, "meta")
haystack/components/converters/image/file_to_document.py
+2 -1
View File
@@ -19,7 +19,8 @@ class ImageFileToDocument:
Converts image file references into empty Document objects with associated metadata.
This component is useful in pipelines where image file paths need to be wrapped in `Document` objects to be
processed by downstream components such as the `SentenceTransformersImageDocumentEmbedder`.
processed by downstream components such as the `LLMDocumentContentExtractor` or the
`SentenceTransformersDocumentImageEmbedder` (available in the `sentence-transformers-haystack` integration).
It does **not** extract any content from the image files, instead it creates `Document` objects with `None` as
their content and attaches metadata such as file path and any user-provided values.
haystack/components/embedders/__init__.py
-16
View File
@@ -12,10 +12,6 @@ _import_structure = {
"azure_text_embedder": ["AzureOpenAITextEmbedder"],
"openai_document_embedder": ["OpenAIDocumentEmbedder"],
"openai_text_embedder": ["OpenAITextEmbedder"],
"sentence_transformers_document_embedder": ["SentenceTransformersDocumentEmbedder"],
"sentence_transformers_text_embedder": ["SentenceTransformersTextEmbedder"],
"sentence_transformers_sparse_document_embedder": ["SentenceTransformersSparseDocumentEmbedder"],
"sentence_transformers_sparse_text_embedder": ["SentenceTransformersSparseTextEmbedder"],
}
if TYPE_CHECKING:
@@ -23,18 +19,6 @@ if TYPE_CHECKING:
from .azure_text_embedder import AzureOpenAITextEmbedder as AzureOpenAITextEmbedder
from .openai_document_embedder import OpenAIDocumentEmbedder as OpenAIDocumentEmbedder
from .openai_text_embedder import OpenAITextEmbedder as OpenAITextEmbedder
from .sentence_transformers_document_embedder import (
SentenceTransformersDocumentEmbedder as SentenceTransformersDocumentEmbedder,
)
from .sentence_transformers_sparse_document_embedder import (
SentenceTransformersSparseDocumentEmbedder as SentenceTransformersSparseDocumentEmbedder,
)
from .sentence_transformers_sparse_text_embedder import (
SentenceTransformersSparseTextEmbedder as SentenceTransformersSparseTextEmbedder,
)
from .sentence_transformers_text_embedder import (
SentenceTransformersTextEmbedder as SentenceTransformersTextEmbedder,
)
else:
sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure)
haystack/components/embedders/backends/__init__.py
-3
View File
@@ -1,3 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
haystack/components/embedders/backends/sentence_transformers_backend.py
-114
View File
@@ -1,114 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
import json
from typing import Any, Literal
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
with LazyImport(message="Run 'pip install \"pillow\"'") as pillow_import:
from PIL.Image import Image
class _SentenceTransformersEmbeddingBackendFactory:
"""
Factory class to create instances of Sentence Transformers embedding backends.
"""
_instances: dict[str, "_SentenceTransformersEmbeddingBackend"] = {}
@staticmethod
def get_embedding_backend(
*,
model: str,
device: str | None = None,
auth_token: Secret | None = None,
trust_remote_code: bool = False,
revision: str | None = None,
local_files_only: bool = False,
truncate_dim: int | None = None,
model_kwargs: dict[str, Any] | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
config_kwargs: dict[str, Any] | None = None,
backend: Literal["torch", "onnx", "openvino"] = "torch",
) -> "_SentenceTransformersEmbeddingBackend":
cache_params = {
"model": model,
"device": device,
"auth_token": auth_token,
"trust_remote_code": trust_remote_code,
"revision": revision,
"local_files_only": local_files_only,
"truncate_dim": truncate_dim,
"model_kwargs": model_kwargs,
"tokenizer_kwargs": tokenizer_kwargs,
"config_kwargs": config_kwargs,
"backend": backend,
}
embedding_backend_id = json.dumps(cache_params, sort_keys=True, default=str)
if embedding_backend_id in _SentenceTransformersEmbeddingBackendFactory._instances:
return _SentenceTransformersEmbeddingBackendFactory._instances[embedding_backend_id]
embedding_backend = _SentenceTransformersEmbeddingBackend(
model=model,
device=device,
auth_token=auth_token,
trust_remote_code=trust_remote_code,
revision=revision,
local_files_only=local_files_only,
truncate_dim=truncate_dim,
model_kwargs=model_kwargs,
tokenizer_kwargs=tokenizer_kwargs,
config_kwargs=config_kwargs,
backend=backend,
)
_SentenceTransformersEmbeddingBackendFactory._instances[embedding_backend_id] = embedding_backend
return embedding_backend
class _SentenceTransformersEmbeddingBackend:
"""
Class to manage Sentence Transformers embeddings.
"""
def __init__(
self,
*,
model: str,
device: str | None = None,
auth_token: Secret | None = None,
trust_remote_code: bool = False,
revision: str | None = None,
local_files_only: bool = False,
truncate_dim: int | None = None,
model_kwargs: dict[str, Any] | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
config_kwargs: dict[str, Any] | None = None,
backend: Literal["torch", "onnx", "openvino"] = "torch",
) -> None:
sentence_transformers_import.check()
self.model = SentenceTransformer(
model_name_or_path=model,
device=device,
token=auth_token.resolve_value() if auth_token else None,
trust_remote_code=trust_remote_code,
revision=revision,
local_files_only=local_files_only,
truncate_dim=truncate_dim,
model_kwargs=model_kwargs,
tokenizer_kwargs=tokenizer_kwargs,
config_kwargs=config_kwargs,
backend=backend,
)
def embed(self, data: list[str] | list["Image"], **kwargs: Any) -> list[list[float]]:
return self.model.encode(data, **kwargs).tolist()
haystack/components/embedders/backends/sentence_transformers_sparse_backend.py
-121
View File
@@ -1,121 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
import json
from typing import Any, Literal
from haystack.dataclasses.sparse_embedding import SparseEmbedding
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 SparseEncoder
class _SentenceTransformersSparseEmbeddingBackendFactory:
"""
Factory class to create instances of Sentence Transformers embedding backends.
"""
_instances: dict[str, "_SentenceTransformersSparseEncoderEmbeddingBackend"] = {}
@staticmethod
def get_embedding_backend(
*,
model: str,
device: str | None = None,
auth_token: Secret | None = None,
trust_remote_code: bool = False,
revision: str | None = None,
local_files_only: bool = False,
model_kwargs: dict[str, Any] | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
config_kwargs: dict[str, Any] | None = None,
backend: Literal["torch", "onnx", "openvino"] = "torch",
) -> "_SentenceTransformersSparseEncoderEmbeddingBackend":
cache_params = {
"model": model,
"device": device,
"auth_token": auth_token,
"trust_remote_code": trust_remote_code,
"revision": revision,
"local_files_only": local_files_only,
"model_kwargs": model_kwargs,
"tokenizer_kwargs": tokenizer_kwargs,
"config_kwargs": config_kwargs,
"backend": backend,
}
embedding_backend_id = json.dumps(cache_params, sort_keys=True, default=str)
if embedding_backend_id in _SentenceTransformersSparseEmbeddingBackendFactory._instances:
return _SentenceTransformersSparseEmbeddingBackendFactory._instances[embedding_backend_id]
embedding_backend = _SentenceTransformersSparseEncoderEmbeddingBackend(
model=model,
device=device,
auth_token=auth_token,
trust_remote_code=trust_remote_code,
revision=revision,
local_files_only=local_files_only,
model_kwargs=model_kwargs,
tokenizer_kwargs=tokenizer_kwargs,
config_kwargs=config_kwargs,
backend=backend,
)
_SentenceTransformersSparseEmbeddingBackendFactory._instances[embedding_backend_id] = embedding_backend
return embedding_backend
class _SentenceTransformersSparseEncoderEmbeddingBackend:
"""
Class to manage Sparse embeddings from Sentence Transformers.
"""
def __init__(
self,
*,
model: str,
device: str | None = None,
auth_token: Secret | None = None,
trust_remote_code: bool = False,
revision: str | None = None,
local_files_only: bool = False,
model_kwargs: dict[str, Any] | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
config_kwargs: dict[str, Any] | None = None,
backend: Literal["torch", "onnx", "openvino"] = "torch",
) -> None:
sentence_transformers_import.check()
self.model = SparseEncoder(
model_name_or_path=model,
device=device,
token=auth_token.resolve_value() if auth_token else None,
trust_remote_code=trust_remote_code,
revision=revision,
local_files_only=local_files_only,
model_kwargs=model_kwargs,
tokenizer_kwargs=tokenizer_kwargs,
config_kwargs=config_kwargs,
backend=backend,
)
def embed(self, *, data: list[str], **kwargs: Any) -> list[SparseEmbedding]:
embeddings_list = self.model.encode(
data,
convert_to_tensor=False, # output is a list of individual tensors
convert_to_sparse_tensor=True,
**kwargs,
)
sparse_embeddings: list[SparseEmbedding] = []
for embedding_tensor in embeddings_list:
embedding_tensor = embedding_tensor.coalesce()
indices = embedding_tensor.indices()[0].tolist() # Only column indices
values = embedding_tensor.values().tolist()
sparse_embeddings.append(SparseEmbedding(indices=indices, values=values))
return sparse_embeddings
haystack/components/embedders/image/__init__.py
-18
View File
@@ -1,18 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
import sys
from typing import TYPE_CHECKING
from lazy_imports import LazyImporter
_import_structure = {"sentence_transformers_doc_image_embedder": ["SentenceTransformersDocumentImageEmbedder"]}
if TYPE_CHECKING:
from .sentence_transformers_doc_image_embedder import (
SentenceTransformersDocumentImageEmbedder as SentenceTransformersDocumentImageEmbedder,
)
else:
sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure)
haystack/components/embedders/image/sentence_transformers_doc_image_embedder.py
-290
View File
@@ -1,290 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from dataclasses import replace
from typing import Any, Literal
from haystack import Document, component, default_from_dict, default_to_dict
from haystack.components.converters.image.image_utils import (
_batch_convert_pdf_pages_to_images,
_extract_image_sources_info,
_PDFPageInfo,
)
from haystack.components.embedders.backends.sentence_transformers_backend import (
_SentenceTransformersEmbeddingBackend,
_SentenceTransformersEmbeddingBackendFactory,
)
from haystack.lazy_imports import LazyImport
from haystack.utils.auth import Secret
from haystack.utils.device import ComponentDevice
from haystack.utils.hf import deserialize_hf_model_kwargs, serialize_hf_model_kwargs
with LazyImport("Run 'pip install pillow'") as pillow_import:
from PIL import Image
@component
class SentenceTransformersDocumentImageEmbedder:
"""
A component for computing Document embeddings based on images using Sentence Transformers models.
The embedding of each Document is stored in the `embedding` field of the Document.
### Usage example
<!-- test-ignore -->
```python
from haystack import Document
from haystack.components.embedders.image import SentenceTransformersDocumentImageEmbedder
embedder = SentenceTransformersDocumentImageEmbedder(model="sentence-transformers/clip-ViT-B-32")
documents = [
Document(content="A photo of a cat", meta={"file_path": "cat.jpg"}),
Document(content="A photo of a dog", meta={"file_path": "dog.jpg"}),
]
result = embedder.run(documents=documents)
documents_with_embeddings = result["documents"]
print(documents_with_embeddings)
# [Document(id=...,
# content='A photo of a cat',
# meta={'file_path': 'cat.jpg',
# 'embedding_source': {'type': 'image', 'file_path_meta_field': 'file_path'}},
# embedding=vector of size 512),
# ...]
```
"""
def __init__( # noqa: PLR0913
self,
*,
file_path_meta_field: str = "file_path",
root_path: str | None = None,
model: str = "sentence-transformers/clip-ViT-B-32",
device: ComponentDevice | None = None,
token: Secret | None = Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False),
batch_size: int = 32,
progress_bar: bool = True,
normalize_embeddings: bool = False,
trust_remote_code: bool = False,
local_files_only: bool = False,
model_kwargs: dict[str, Any] | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
config_kwargs: dict[str, Any] | None = None,
precision: Literal["float32", "int8", "uint8", "binary", "ubinary"] = "float32",
encode_kwargs: dict[str, Any] | None = None,
backend: Literal["torch", "onnx", "openvino"] = "torch",
) -> None:
"""
Creates a SentenceTransformersDocumentEmbedder component.
:param file_path_meta_field: The metadata field in the Document that contains the file path to the image or PDF.
:param root_path: The root directory path where document files are located. If provided, file paths in
document metadata will be resolved relative to this path. If None, file paths are treated as absolute paths.
:param model:
The Sentence Transformers model to use for calculating embeddings. Pass a local path or ID of the model on
Hugging Face. To be used with this component, the model must be able to embed images and text into the same
vector space. Compatible models include:
- "sentence-transformers/clip-ViT-B-32"
- "sentence-transformers/clip-ViT-L-14"
- "sentence-transformers/clip-ViT-B-16"
- "sentence-transformers/clip-ViT-B-32-multilingual-v1"
- "jinaai/jina-embeddings-v4"
- "jinaai/jina-clip-v1"
- "jinaai/jina-clip-v2".
:param device:
The device to use for loading the model.
Overrides the default device.
:param token:
The API token to download private models from Hugging Face.
:param batch_size:
Number of documents to embed at once.
:param progress_bar:
If `True`, shows a progress bar when embedding documents.
:param normalize_embeddings:
If `True`, the embeddings are normalized using L2 normalization, so that each embedding has a norm of 1.
:param trust_remote_code:
If `False`, allows only Hugging Face verified model architectures.
If `True`, allows custom models and scripts.
:param local_files_only:
If `True`, does not attempt to download the model from Hugging Face Hub and only looks at local files.
:param model_kwargs:
Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
when loading the model. Refer to specific model documentation for available kwargs.
:param tokenizer_kwargs:
Additional keyword arguments for `AutoTokenizer.from_pretrained` when loading the tokenizer.
Refer to specific model documentation for available kwargs.
:param config_kwargs:
Additional keyword arguments for `AutoConfig.from_pretrained` when loading the model configuration.
:param precision:
The precision to use for the embeddings.
All non-float32 precisions are quantized embeddings.
Quantized embeddings are smaller and faster to compute, but may have a lower accuracy.
They are useful for reducing the size of the embeddings of a corpus for semantic search, among other tasks.
:param encode_kwargs:
Additional keyword arguments for `SentenceTransformer.encode` when embedding documents.
This parameter is provided for fine customization. Be careful not to clash with already set parameters and
avoid passing parameters that change the output type.
:param backend:
The backend to use for the Sentence Transformers model. Choose from "torch", "onnx", or "openvino".
Refer to the [Sentence Transformers documentation](https://sbert.net/docs/sentence_transformer/usage/efficiency.html)
for more information on acceleration and quantization options.
"""
pillow_import.check()
self.file_path_meta_field = file_path_meta_field
self.root_path = root_path or ""
self.model = model
self.device = ComponentDevice.resolve_device(device)
self.token = token
self.batch_size = batch_size
self.progress_bar = progress_bar
self.normalize_embeddings = normalize_embeddings
self.trust_remote_code = trust_remote_code
self.local_files_only = local_files_only
self.model_kwargs = model_kwargs
self.tokenizer_kwargs = tokenizer_kwargs
self.config_kwargs = config_kwargs
self.encode_kwargs = encode_kwargs
self.precision = precision
self.backend = backend
self._embedding_backend: _SentenceTransformersEmbeddingBackend | None = None
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.
:returns:
Dictionary with serialized data.
"""
serialization_dict = default_to_dict(
self,
file_path_meta_field=self.file_path_meta_field,
root_path=self.root_path,
model=self.model,
device=self.device,
token=self.token,
batch_size=self.batch_size,
progress_bar=self.progress_bar,
normalize_embeddings=self.normalize_embeddings,
trust_remote_code=self.trust_remote_code,
local_files_only=self.local_files_only,
model_kwargs=self.model_kwargs,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
precision=self.precision,
encode_kwargs=self.encode_kwargs,
backend=self.backend,
)
if serialization_dict["init_parameters"].get("model_kwargs") is not None:
serialize_hf_model_kwargs(serialization_dict["init_parameters"]["model_kwargs"])
return serialization_dict
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "SentenceTransformersDocumentImageEmbedder":
"""
Deserializes the component from a dictionary.
:param data:
Dictionary to deserialize from.
:returns:
Deserialized component.
"""
init_params = data["init_parameters"]
if init_params.get("model_kwargs") is not None:
deserialize_hf_model_kwargs(init_params["model_kwargs"])
return default_from_dict(cls, data)
def warm_up(self) -> None:
"""
Initializes the component.
"""
if self._embedding_backend is None:
self._embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model=self.model,
device=self.device.to_torch_str(),
auth_token=self.token,
trust_remote_code=self.trust_remote_code,
local_files_only=self.local_files_only,
model_kwargs=self.model_kwargs,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
backend=self.backend,
)
if self.tokenizer_kwargs and self.tokenizer_kwargs.get("model_max_length"):
self._embedding_backend.model.max_seq_length = self.tokenizer_kwargs["model_max_length"]
@component.output_types(documents=list[Document])
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
"""
Embed a list of documents.
:param documents:
Documents to embed.
:returns:
A dictionary with the following keys:
- `documents`: Documents with embeddings.
"""
if not isinstance(documents, list) or documents and not isinstance(documents[0], Document):
raise TypeError(
"SentenceTransformersDocumentImageEmbedder expects a list of Documents as input. "
"In case you want to embed a string, please use the SentenceTransformersTextEmbedder."
)
if self._embedding_backend is None:
self.warm_up()
images_source_info = _extract_image_sources_info(
documents=documents, file_path_meta_field=self.file_path_meta_field, root_path=self.root_path
)
images_to_embed: list = [None] * len(documents)
pdf_page_infos: list[_PDFPageInfo] = []
for doc_idx, image_source_info in enumerate(images_source_info):
if image_source_info["mime_type"] == "application/pdf":
# Store PDF documents for later processing
page_number = image_source_info.get("page_number")
assert page_number is not None # checked in _extract_image_sources_info but mypy doesn't know that
pdf_page_info: _PDFPageInfo = {
"doc_idx": doc_idx,
"path": image_source_info["path"],
"page_number": page_number,
}
pdf_page_infos.append(pdf_page_info)
else:
# Process images directly
image = Image.open(image_source_info["path"])
images_to_embed[doc_idx] = image
pdf_images_by_doc_idx = _batch_convert_pdf_pages_to_images(pdf_page_infos=pdf_page_infos, return_base64=False)
for doc_idx, pil_image in pdf_images_by_doc_idx.items():
images_to_embed[doc_idx] = pil_image
none_images_doc_ids = [documents[doc_idx].id for doc_idx, image in enumerate(images_to_embed) if image is None]
if none_images_doc_ids:
raise RuntimeError(f"Conversion failed for some documents. Document IDs: {none_images_doc_ids}.")
# mypy doesn't know this is set in warm_up
embeddings = self._embedding_backend.embed( # type: ignore[union-attr]
data=images_to_embed,
batch_size=self.batch_size,
show_progress_bar=self.progress_bar,
normalize_embeddings=self.normalize_embeddings,
precision=self.precision,
**(self.encode_kwargs if self.encode_kwargs else {}),
)
docs_with_embeddings = []
for doc, emb in zip(documents, embeddings, strict=True):
# we store this information for later inspection
new_meta = {
**doc.meta,
"embedding_source": {"type": "image", "file_path_meta_field": self.file_path_meta_field},
}
new_doc = replace(doc, meta=new_meta, embedding=emb)
docs_with_embeddings.append(new_doc)
return {"documents": docs_with_embeddings}
haystack/components/embedders/sentence_transformers_document_embedder.py
-269
View File
@@ -1,269 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from dataclasses import replace
from typing import Any, Literal
from haystack import Document, component, default_from_dict, default_to_dict
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
@component
class SentenceTransformersDocumentEmbedder:
"""
Calculates document embeddings using Sentence Transformers models.
It stores the embeddings in the `embedding` metadata field of each document.
You can also embed documents' metadata.
Use this component in indexing pipelines to embed input documents
and send them to DocumentWriter to write into a Document Store.
### Usage example:
<!-- test-ignore -->
```python
from haystack import Document
from haystack.components.embedders import SentenceTransformersDocumentEmbedder
doc = Document(content="I love pizza!")
doc_embedder = SentenceTransformersDocumentEmbedder()
result = doc_embedder.run([doc])
print(result['documents'][0].embedding)
# [-0.07804739475250244, 0.1498992145061493, ...]
```
"""
def __init__( # noqa: PLR0913
self,
model: str = "sentence-transformers/all-mpnet-base-v2",
device: ComponentDevice | None = None,
token: Secret | None = Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False),
prefix: str = "",
suffix: str = "",
batch_size: int = 32,
progress_bar: bool = True,
normalize_embeddings: bool = False,
meta_fields_to_embed: list[str] | None = None,
embedding_separator: str = "\n",
trust_remote_code: bool = False,
local_files_only: bool = False,
truncate_dim: int | None = None,
model_kwargs: dict[str, Any] | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
config_kwargs: dict[str, Any] | None = None,
precision: Literal["float32", "int8", "uint8", "binary", "ubinary"] = "float32",
encode_kwargs: dict[str, Any] | None = None,
backend: Literal["torch", "onnx", "openvino"] = "torch",
revision: str | None = None,
) -> None:
"""
Creates a SentenceTransformersDocumentEmbedder component.
:param model:
The model to use for calculating embeddings.
Pass a local path or ID of the model on Hugging Face.
:param device:
The device to use for loading the model.
Overrides the default device.
:param token:
The API token to download private models from Hugging Face.
:param prefix:
A string to add at the beginning of each document text.
Can be used to prepend the text with an instruction, as required by some embedding models,
such as E5 and bge.
:param suffix:
A string to add at the end of each document text.
:param batch_size:
Number of documents to embed at once.
:param progress_bar:
If `True`, shows a progress bar when embedding documents.
:param normalize_embeddings:
If `True`, the embeddings are normalized using L2 normalization, so that each embedding has a norm of 1.
:param meta_fields_to_embed:
List of metadata fields to embed along with the document text.
:param embedding_separator:
Separator used to concatenate the metadata fields to the document text.
:param trust_remote_code:
If `False`, allows only Hugging Face verified model architectures.
If `True`, allows custom models and scripts.
:param local_files_only:
If `True`, does not attempt to download the model from Hugging Face Hub and only looks at local files.
:param truncate_dim:
The dimension to truncate sentence embeddings to. `None` does no truncation.
If the model wasn't trained with Matryoshka Representation Learning,
truncating embeddings can significantly affect performance.
:param model_kwargs:
Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
when loading the model. Refer to specific model documentation for available kwargs.
:param tokenizer_kwargs:
Additional keyword arguments for `AutoTokenizer.from_pretrained` when loading the tokenizer.
Refer to specific model documentation for available kwargs.
:param config_kwargs:
Additional keyword arguments for `AutoConfig.from_pretrained` when loading the model configuration.
:param precision:
The precision to use for the embeddings.
All non-float32 precisions are quantized embeddings.
Quantized embeddings are smaller and faster to compute, but may have a lower accuracy.
They are useful for reducing the size of the embeddings of a corpus for semantic search, among other tasks.
:param encode_kwargs:
Additional keyword arguments for `SentenceTransformer.encode` when embedding documents.
This parameter is provided for fine customization. Be careful not to clash with already set parameters and
avoid passing parameters that change the output type.
:param backend:
The backend to use for the Sentence Transformers model. Choose from "torch", "onnx", or "openvino".
Refer to the [Sentence Transformers documentation](https://sbert.net/docs/sentence_transformer/usage/efficiency.html)
for more information on acceleration and quantization options.
: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.
"""
self.model = model
self.device = ComponentDevice.resolve_device(device)
self.token = token
self.prefix = prefix
self.suffix = suffix
self.batch_size = batch_size
self.progress_bar = progress_bar
self.normalize_embeddings = normalize_embeddings
self.meta_fields_to_embed = meta_fields_to_embed or []
self.embedding_separator = embedding_separator
self.trust_remote_code = trust_remote_code
self.revision = revision
self.local_files_only = local_files_only
self.truncate_dim = truncate_dim
self.model_kwargs = model_kwargs
self.tokenizer_kwargs = tokenizer_kwargs
self.config_kwargs = config_kwargs
self.encode_kwargs = encode_kwargs
self.embedding_backend: _SentenceTransformersEmbeddingBackend | None = None
self.precision = precision
self.backend = backend
def _get_telemetry_data(self) -> dict[str, Any]:
"""
Data that is sent to Posthog for usage analytics.
"""
return {"model": self.model}
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.
:returns:
Dictionary with serialized data.
"""
serialization_dict = default_to_dict(
self,
model=self.model,
device=self.device,
token=self.token,
prefix=self.prefix,
suffix=self.suffix,
batch_size=self.batch_size,
progress_bar=self.progress_bar,
normalize_embeddings=self.normalize_embeddings,
meta_fields_to_embed=self.meta_fields_to_embed,
embedding_separator=self.embedding_separator,
trust_remote_code=self.trust_remote_code,
revision=self.revision,
local_files_only=self.local_files_only,
truncate_dim=self.truncate_dim,
model_kwargs=self.model_kwargs,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
precision=self.precision,
encode_kwargs=self.encode_kwargs,
backend=self.backend,
)
if serialization_dict["init_parameters"].get("model_kwargs") is not None:
serialize_hf_model_kwargs(serialization_dict["init_parameters"]["model_kwargs"])
return serialization_dict
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "SentenceTransformersDocumentEmbedder":
"""
Deserializes the component from a dictionary.
:param data:
Dictionary to deserialize from.
:returns:
Deserialized component.
"""
init_params = data["init_parameters"]
if init_params.get("model_kwargs") is not None:
deserialize_hf_model_kwargs(init_params["model_kwargs"])
return default_from_dict(cls, data)
def warm_up(self) -> None:
"""
Initializes the component.
"""
if self.embedding_backend is None:
self.embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model=self.model,
device=self.device.to_torch_str(),
auth_token=self.token,
trust_remote_code=self.trust_remote_code,
revision=self.revision,
local_files_only=self.local_files_only,
truncate_dim=self.truncate_dim,
model_kwargs=self.model_kwargs,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
backend=self.backend,
)
if self.tokenizer_kwargs and self.tokenizer_kwargs.get("model_max_length"):
self.embedding_backend.model.max_seq_length = self.tokenizer_kwargs["model_max_length"]
@component.output_types(documents=list[Document])
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
"""
Embed a list of documents.
:param documents:
Documents to embed.
:returns:
A dictionary with the following keys:
- `documents`: Documents with embeddings.
"""
if not isinstance(documents, list) or documents and not isinstance(documents[0], Document):
raise TypeError(
"SentenceTransformersDocumentEmbedder expects a list of Documents as input."
"In case you want to embed a string, please use the SentenceTransformersTextEmbedder."
)
if self.embedding_backend is None:
self.warm_up()
texts_to_embed = []
for doc in documents:
meta_values_to_embed = [
str(doc.meta[key]) for key in self.meta_fields_to_embed if key in doc.meta and doc.meta[key] is not None
]
text_to_embed = (
self.prefix + self.embedding_separator.join(meta_values_to_embed + [doc.content or ""]) + self.suffix
)
texts_to_embed.append(text_to_embed)
# # mypy doesn't know this is set in warm_up
embeddings = self.embedding_backend.embed( # type: ignore[union-attr]
texts_to_embed,
batch_size=self.batch_size,
show_progress_bar=self.progress_bar,
normalize_embeddings=self.normalize_embeddings,
precision=self.precision,
**(self.encode_kwargs if self.encode_kwargs else {}),
)
new_documents = []
for doc, emb in zip(documents, embeddings, strict=True):
new_documents.append(replace(doc, embedding=emb))
return {"documents": new_documents}
haystack/components/embedders/sentence_transformers_sparse_document_embedder.py
-236
View File
@@ -1,236 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from dataclasses import replace
from typing import Any, Literal
from haystack import Document, component, default_from_dict, default_to_dict
from haystack.components.embedders.backends.sentence_transformers_sparse_backend import (
_SentenceTransformersSparseEmbeddingBackendFactory,
_SentenceTransformersSparseEncoderEmbeddingBackend,
)
from haystack.utils import ComponentDevice, Secret
from haystack.utils.hf import deserialize_hf_model_kwargs, serialize_hf_model_kwargs
@component
class SentenceTransformersSparseDocumentEmbedder:
"""
Calculates document sparse embeddings using sparse embedding models from Sentence Transformers.
It stores the sparse embeddings in the `sparse_embedding` metadata field of each document.
You can also embed documents' metadata.
Use this component in indexing pipelines to embed input documents
and send them to DocumentWriter to write a into a Document Store.
### Usage example:
<!-- test-ignore -->
```python
from haystack import Document
from haystack.components.embedders import SentenceTransformersSparseDocumentEmbedder
doc = Document(content="I love pizza!")
doc_embedder = SentenceTransformersSparseDocumentEmbedder()
result = doc_embedder.run([doc])
print(result['documents'][0].sparse_embedding)
# SparseEmbedding(indices=[999, 1045, ...], values=[0.918, 0.867, ...])
```
"""
def __init__( # noqa: PLR0913
self,
*,
model: str = "prithivida/Splade_PP_en_v2",
device: ComponentDevice | None = None,
token: Secret | None = Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False),
prefix: str = "",
suffix: str = "",
batch_size: int = 32,
progress_bar: bool = True,
meta_fields_to_embed: list[str] | None = None,
embedding_separator: str = "\n",
trust_remote_code: bool = False,
local_files_only: bool = False,
model_kwargs: dict[str, Any] | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
config_kwargs: dict[str, Any] | None = None,
backend: Literal["torch", "onnx", "openvino"] = "torch",
revision: str | None = None,
) -> None:
"""
Creates a SentenceTransformersSparseDocumentEmbedder component.
:param model:
The model to use for calculating sparse embeddings.
Pass a local path or ID of the model on Hugging Face.
:param device:
The device to use for loading the model.
Overrides the default device.
:param token:
The API token to download private models from Hugging Face.
:param prefix:
A string to add at the beginning of each document text.
:param suffix:
A string to add at the end of each document text.
:param batch_size:
Number of documents to embed at once.
:param progress_bar:
If `True`, shows a progress bar when embedding documents.
:param meta_fields_to_embed:
List of metadata fields to embed along with the document text.
:param embedding_separator:
Separator used to concatenate the metadata fields to the document text.
:param trust_remote_code:
If `False`, allows only Hugging Face verified model architectures.
If `True`, allows custom models and scripts.
:param local_files_only:
If `True`, does not attempt to download the model from Hugging Face Hub and only looks at local files.
:param model_kwargs:
Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
when loading the model. Refer to specific model documentation for available kwargs.
:param tokenizer_kwargs:
Additional keyword arguments for `AutoTokenizer.from_pretrained` when loading the tokenizer.
Refer to specific model documentation for available kwargs.
:param config_kwargs:
Additional keyword arguments for `AutoConfig.from_pretrained` when loading the model configuration.
:param backend:
The backend to use for the Sentence Transformers model. Choose from "torch", "onnx", or "openvino".
Refer to the [Sentence Transformers documentation](https://sbert.net/docs/sentence_transformer/usage/efficiency.html)
for more information on acceleration and quantization options.
: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.
"""
self.model = model
self.device = ComponentDevice.resolve_device(device)
self.token = token
self.prefix = prefix
self.suffix = suffix
self.batch_size = batch_size
self.progress_bar = progress_bar
self.meta_fields_to_embed = meta_fields_to_embed or []
self.embedding_separator = embedding_separator
self.trust_remote_code = trust_remote_code
self.revision = revision
self.local_files_only = local_files_only
self.model_kwargs = model_kwargs
self.tokenizer_kwargs = tokenizer_kwargs
self.config_kwargs = config_kwargs
self.embedding_backend: _SentenceTransformersSparseEncoderEmbeddingBackend | None = None
self.backend = backend
def _get_telemetry_data(self) -> dict[str, Any]:
"""
Data that is sent to Posthog for usage analytics.
"""
return {"model": self.model}
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.
:returns:
Dictionary with serialized data.
"""
serialization_dict = default_to_dict(
self,
model=self.model,
device=self.device,
token=self.token,
prefix=self.prefix,
suffix=self.suffix,
batch_size=self.batch_size,
progress_bar=self.progress_bar,
meta_fields_to_embed=self.meta_fields_to_embed,
embedding_separator=self.embedding_separator,
trust_remote_code=self.trust_remote_code,
revision=self.revision,
local_files_only=self.local_files_only,
model_kwargs=self.model_kwargs,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
backend=self.backend,
)
if serialization_dict["init_parameters"].get("model_kwargs") is not None:
serialize_hf_model_kwargs(serialization_dict["init_parameters"]["model_kwargs"])
return serialization_dict
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "SentenceTransformersSparseDocumentEmbedder":
"""
Deserializes the component from a dictionary.
:param data:
Dictionary to deserialize from.
:returns:
Deserialized component.
"""
init_params = data["init_parameters"]
if init_params.get("model_kwargs") is not None:
deserialize_hf_model_kwargs(init_params["model_kwargs"])
return default_from_dict(cls, data)
def warm_up(self) -> None:
"""
Initializes the component.
"""
if self.embedding_backend is None:
self.embedding_backend = _SentenceTransformersSparseEmbeddingBackendFactory.get_embedding_backend(
model=self.model,
device=self.device.to_torch_str(),
auth_token=self.token,
trust_remote_code=self.trust_remote_code,
revision=self.revision,
local_files_only=self.local_files_only,
model_kwargs=self.model_kwargs,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
backend=self.backend,
)
if self.tokenizer_kwargs and self.tokenizer_kwargs.get("model_max_length"):
self.embedding_backend.model.max_seq_length = self.tokenizer_kwargs["model_max_length"]
@component.output_types(documents=list[Document])
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
"""
Embed a list of documents.
:param documents:
Documents to embed.
:returns:
A dictionary with the following keys:
- `documents`: Documents with sparse embeddings under the `sparse_embedding` field.
"""
if not isinstance(documents, list) or documents and not isinstance(documents[0], Document):
raise TypeError(
"SentenceTransformersSparseDocumentEmbedder expects a list of Documents as input."
"In case you want to embed a list of strings, please use the SentenceTransformersSparseTextEmbedder."
)
if self.embedding_backend is None:
self.warm_up()
texts_to_embed = []
for doc in documents:
meta_values_to_embed = [
str(doc.meta[key]) for key in self.meta_fields_to_embed if key in doc.meta and doc.meta[key] is not None
]
text_to_embed = (
self.prefix + self.embedding_separator.join(meta_values_to_embed + [doc.content or ""]) + self.suffix
)
texts_to_embed.append(text_to_embed)
# mypy doesn't know this is set in warm_up
embeddings = self.embedding_backend.embed( # type: ignore[union-attr]
data=texts_to_embed, batch_size=self.batch_size, show_progress_bar=self.progress_bar
)
documents_with_embeddings = []
for doc, emb in zip(documents, embeddings, strict=True):
documents_with_embeddings.append(replace(doc, sparse_embedding=emb))
return {"documents": documents_with_embeddings}
haystack/components/embedders/sentence_transformers_sparse_text_embedder.py
-198
View File
@@ -1,198 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from typing import Any, Literal
from haystack import component, default_from_dict, default_to_dict
from haystack.components.embedders.backends.sentence_transformers_sparse_backend import (
_SentenceTransformersSparseEmbeddingBackendFactory,
_SentenceTransformersSparseEncoderEmbeddingBackend,
)
from haystack.dataclasses.sparse_embedding import SparseEmbedding
from haystack.utils import ComponentDevice, Secret
from haystack.utils.hf import deserialize_hf_model_kwargs, serialize_hf_model_kwargs
@component
class SentenceTransformersSparseTextEmbedder:
"""
Embeds strings using sparse embedding models from Sentence Transformers.
You can use it to embed user query and send it to a sparse embedding retriever.
Usage example:
<!-- test-ignore -->
```python
from haystack.components.embedders import SentenceTransformersSparseTextEmbedder
text_to_embed = "I love pizza!"
text_embedder = SentenceTransformersSparseTextEmbedder()
print(text_embedder.run(text_to_embed))
# {'sparse_embedding': SparseEmbedding(indices=[999, 1045, ...], values=[0.918, 0.867, ...])}
```
"""
def __init__( # noqa: PLR0913
self,
*,
model: str = "prithivida/Splade_PP_en_v2",
device: ComponentDevice | None = None,
token: Secret | None = Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False),
prefix: str = "",
suffix: str = "",
trust_remote_code: bool = False,
local_files_only: bool = False,
model_kwargs: dict[str, Any] | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
config_kwargs: dict[str, Any] | None = None,
backend: Literal["torch", "onnx", "openvino"] = "torch",
revision: str | None = None,
) -> None:
"""
Create a SentenceTransformersSparseTextEmbedder component.
:param model:
The model to use for calculating sparse embeddings.
Specify the path to a local model or the ID of the model on Hugging Face.
:param device:
Overrides the default device used to load the model.
:param token:
An API token to use private models from Hugging Face.
:param prefix:
A string to add at the beginning of each text to be embedded.
:param suffix:
A string to add at the end of each text to embed.
:param trust_remote_code:
If `False`, permits only Hugging Face verified model architectures.
If `True`, permits custom models and scripts.
:param local_files_only:
If `True`, does not attempt to download the model from Hugging Face Hub and only looks at local files.
:param model_kwargs:
Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
when loading the model. Refer to specific model documentation for available kwargs.
:param tokenizer_kwargs:
Additional keyword arguments for `AutoTokenizer.from_pretrained` when loading the tokenizer.
Refer to specific model documentation for available kwargs.
:param config_kwargs:
Additional keyword arguments for `AutoConfig.from_pretrained` when loading the model configuration.
:param backend:
The backend to use for the Sentence Transformers model. Choose from "torch", "onnx", or "openvino".
Refer to the [Sentence Transformers documentation](https://sbert.net/docs/sentence_transformer/usage/efficiency.html)
for more information on acceleration and quantization options.
: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.
"""
self.model = model
self.device = ComponentDevice.resolve_device(device)
self.token = token
self.prefix = prefix
self.suffix = suffix
self.trust_remote_code = trust_remote_code
self.revision = revision
self.local_files_only = local_files_only
self.model_kwargs = model_kwargs
self.tokenizer_kwargs = tokenizer_kwargs
self.config_kwargs = config_kwargs
self.embedding_backend: _SentenceTransformersSparseEncoderEmbeddingBackend | None = None
self.backend = backend
def _get_telemetry_data(self) -> dict[str, Any]:
"""
Data that is sent to Posthog for usage analytics.
"""
return {"model": self.model}
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.
:returns:
Dictionary with serialized data.
"""
serialization_dict = default_to_dict(
self,
model=self.model,
device=self.device,
token=self.token,
prefix=self.prefix,
suffix=self.suffix,
trust_remote_code=self.trust_remote_code,
revision=self.revision,
local_files_only=self.local_files_only,
model_kwargs=self.model_kwargs,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
backend=self.backend,
)
if serialization_dict["init_parameters"].get("model_kwargs") is not None:
serialize_hf_model_kwargs(serialization_dict["init_parameters"]["model_kwargs"])
return serialization_dict
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "SentenceTransformersSparseTextEmbedder":
"""
Deserializes the component from a dictionary.
:param data:
Dictionary to deserialize from.
:returns:
Deserialized component.
"""
init_params = data["init_parameters"]
if init_params.get("model_kwargs") is not None:
deserialize_hf_model_kwargs(init_params["model_kwargs"])
return default_from_dict(cls, data)
def warm_up(self) -> None:
"""
Initializes the component.
"""
if self.embedding_backend is None:
self.embedding_backend = _SentenceTransformersSparseEmbeddingBackendFactory.get_embedding_backend(
model=self.model,
device=self.device.to_torch_str(),
auth_token=self.token,
trust_remote_code=self.trust_remote_code,
revision=self.revision,
local_files_only=self.local_files_only,
model_kwargs=self.model_kwargs,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
backend=self.backend,
)
if self.tokenizer_kwargs and self.tokenizer_kwargs.get("model_max_length"):
self.embedding_backend.model.max_seq_length = self.tokenizer_kwargs["model_max_length"]
@component.output_types(sparse_embedding=SparseEmbedding)
def run(self, text: str) -> dict[str, Any]:
"""
Embed a single string.
:param text:
Text to embed.
:returns:
A dictionary with the following keys:
- `sparse_embedding`: The sparse embedding of the input text.
"""
if not isinstance(text, str):
raise TypeError(
"SentenceTransformersSparseTextEmbedder expects a string as input."
"In case you want to embed a list of Documents, please use the"
"SentenceTransformersSparseDocumentEmbedder."
)
if self.embedding_backend is None:
self.warm_up()
text_to_embed = self.prefix + text + self.suffix
# mypy doesn't know this is set in warm_up
sparse_embedding = self.embedding_backend.embed(data=[text_to_embed])[0] # type: ignore[union-attr]
return {"sparse_embedding": sparse_embedding}
haystack/components/embedders/sentence_transformers_text_embedder.py
-242
View File
@@ -1,242 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from typing import Any, Literal
from haystack import component, default_from_dict, default_to_dict
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
@component
class SentenceTransformersTextEmbedder:
"""
Embeds strings using Sentence Transformers models.
You can use it to embed user query and send it to an embedding retriever.
Usage example:
<!-- test-ignore -->
```python
from haystack.components.embedders import SentenceTransformersTextEmbedder
text_to_embed = "I love pizza!"
text_embedder = SentenceTransformersTextEmbedder()
print(text_embedder.run(text_to_embed))
# {'embedding': [-0.07804739475250244, 0.1498992145061493,, ...]}
```
"""
def __init__( # noqa: PLR0913
self,
model: str = "sentence-transformers/all-mpnet-base-v2",
device: ComponentDevice | None = None,
token: Secret | None = Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False),
prefix: str = "",
suffix: str = "",
batch_size: int = 32,
progress_bar: bool = True,
normalize_embeddings: bool = False,
trust_remote_code: bool = False,
local_files_only: bool = False,
truncate_dim: int | None = None,
model_kwargs: dict[str, Any] | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
config_kwargs: dict[str, Any] | None = None,
precision: Literal["float32", "int8", "uint8", "binary", "ubinary"] = "float32",
encode_kwargs: dict[str, Any] | None = None,
backend: Literal["torch", "onnx", "openvino"] = "torch",
revision: str | None = None,
) -> None:
"""
Create a SentenceTransformersTextEmbedder component.
:param model:
The model to use for calculating embeddings.
Specify the path to a local model or the ID of the model on Hugging Face.
:param device:
Overrides the default device used to load the model.
:param token:
An API token to use private models from Hugging Face.
:param prefix:
A string to add at the beginning of each text to be embedded.
You can use it to prepend the text with an instruction, as required by some embedding models,
such as E5 and bge.
:param suffix:
A string to add at the end of each text to embed.
:param batch_size:
Number of texts to embed at once.
:param progress_bar:
If `True`, shows a progress bar for calculating embeddings.
If `False`, disables the progress bar.
:param normalize_embeddings:
If `True`, the embeddings are normalized using L2 normalization, so that the embeddings have a norm of 1.
:param trust_remote_code:
If `False`, permits only Hugging Face verified model architectures.
If `True`, permits custom models and scripts.
:param local_files_only:
If `True`, does not attempt to download the model from Hugging Face Hub and only looks at local files.
:param truncate_dim:
The dimension to truncate sentence embeddings to. `None` does no truncation.
If the model has not been trained with Matryoshka Representation Learning,
truncation of embeddings can significantly affect performance.
:param model_kwargs:
Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
when loading the model. Refer to specific model documentation for available kwargs.
:param tokenizer_kwargs:
Additional keyword arguments for `AutoTokenizer.from_pretrained` when loading the tokenizer.
Refer to specific model documentation for available kwargs.
:param config_kwargs:
Additional keyword arguments for `AutoConfig.from_pretrained` when loading the model configuration.
:param precision:
The precision to use for the embeddings.
All non-float32 precisions are quantized embeddings.
Quantized embeddings are smaller in size and faster to compute, but may have a lower accuracy.
They are useful for reducing the size of the embeddings of a corpus for semantic search, among other tasks.
:param encode_kwargs:
Additional keyword arguments for `SentenceTransformer.encode` when embedding texts.
This parameter is provided for fine customization. Be careful not to clash with already set parameters and
avoid passing parameters that change the output type.
:param backend:
The backend to use for the Sentence Transformers model. Choose from "torch", "onnx", or "openvino".
Refer to the [Sentence Transformers documentation](https://sbert.net/docs/sentence_transformer/usage/efficiency.html)
for more information on acceleration and quantization options.
: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.
"""
self.model = model
self.device = ComponentDevice.resolve_device(device)
self.token = token
self.prefix = prefix
self.suffix = suffix
self.batch_size = batch_size
self.progress_bar = progress_bar
self.normalize_embeddings = normalize_embeddings
self.trust_remote_code = trust_remote_code
self.revision = revision
self.local_files_only = local_files_only
self.truncate_dim = truncate_dim
self.model_kwargs = model_kwargs
self.tokenizer_kwargs = tokenizer_kwargs
self.config_kwargs = config_kwargs
self.encode_kwargs = encode_kwargs
self.embedding_backend: _SentenceTransformersEmbeddingBackend | None = None
self.precision = precision
self.backend = backend
def _get_telemetry_data(self) -> dict[str, Any]:
"""
Data that is sent to Posthog for usage analytics.
"""
return {"model": self.model}
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.
:returns:
Dictionary with serialized data.
"""
serialization_dict = default_to_dict(
self,
model=self.model,
device=self.device,
token=self.token,
prefix=self.prefix,
suffix=self.suffix,
batch_size=self.batch_size,
progress_bar=self.progress_bar,
normalize_embeddings=self.normalize_embeddings,
trust_remote_code=self.trust_remote_code,
revision=self.revision,
local_files_only=self.local_files_only,
truncate_dim=self.truncate_dim,
model_kwargs=self.model_kwargs,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
precision=self.precision,
encode_kwargs=self.encode_kwargs,
backend=self.backend,
)
if serialization_dict["init_parameters"].get("model_kwargs") is not None:
serialize_hf_model_kwargs(serialization_dict["init_parameters"]["model_kwargs"])
return serialization_dict
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "SentenceTransformersTextEmbedder":
"""
Deserializes the component from a dictionary.
:param data:
Dictionary to deserialize from.
:returns:
Deserialized component.
"""
init_params = data["init_parameters"]
if init_params.get("model_kwargs") is not None:
deserialize_hf_model_kwargs(init_params["model_kwargs"])
return default_from_dict(cls, data)
def warm_up(self) -> None:
"""
Initializes the component.
"""
if self.embedding_backend is None:
self.embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model=self.model,
device=self.device.to_torch_str(),
auth_token=self.token,
trust_remote_code=self.trust_remote_code,
revision=self.revision,
local_files_only=self.local_files_only,
truncate_dim=self.truncate_dim,
model_kwargs=self.model_kwargs,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
backend=self.backend,
)
if self.tokenizer_kwargs and self.tokenizer_kwargs.get("model_max_length"):
self.embedding_backend.model.max_seq_length = self.tokenizer_kwargs["model_max_length"]
@component.output_types(embedding=list[float])
def run(self, text: str) -> dict[str, Any]:
"""
Embed a single string.
:param text:
Text to embed.
:returns:
A dictionary with the following keys:
- `embedding`: The embedding of the input text.
"""
if not isinstance(text, str):
raise TypeError(
"SentenceTransformersTextEmbedder expects a string as input."
"In case you want to embed a list of Documents, please use the SentenceTransformersDocumentEmbedder."
)
if self.embedding_backend is None:
self.warm_up()
text_to_embed = self.prefix + text + self.suffix
# mypy doesn't know this is set in warm_up
embedding = self.embedding_backend.embed( # type: ignore[union-attr]
[text_to_embed],
batch_size=self.batch_size,
show_progress_bar=self.progress_bar,
normalize_embeddings=self.normalize_embeddings,
precision=self.precision,
**(self.encode_kwargs if self.encode_kwargs else {}),
)[0]
return {"embedding": embedding}
haystack/components/joiners/document_joiner.py
+3 -3
View File
@@ -57,7 +57,7 @@ class DocumentJoiner:
```python
from haystack import Pipeline, Document
from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder
from haystack.components.embedders import OpenAITextEmbedder, OpenAIDocumentEmbedder
from haystack.components.joiners import DocumentJoiner
from haystack.components.retrievers import InMemoryBM25Retriever
from haystack.components.retrievers import InMemoryEmbeddingRetriever
@@ -65,14 +65,14 @@ class DocumentJoiner:
document_store = InMemoryDocumentStore()
docs = [Document(content="Paris"), Document(content="Berlin"), Document(content="London")]
embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
embedder = OpenAIDocumentEmbedder()
docs_embeddings = embedder.run(docs)
document_store.write_documents(docs_embeddings['documents'])
p = Pipeline()
p.add_component(instance=InMemoryBM25Retriever(document_store=document_store), name="bm25_retriever")
p.add_component(
instance=SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"),
instance=OpenAITextEmbedder(),
name="text_embedder",
)
p.add_component(instance=InMemoryEmbeddingRetriever(document_store=document_store), name="embedding_retriever")
haystack/components/preprocessors/embedding_based_document_splitter.py
+2 -2
View File
@@ -38,7 +38,7 @@ class EmbeddingBasedDocumentSplitter:
```python
from haystack import Document
from haystack.components.embedders import SentenceTransformersDocumentEmbedder
from haystack.components.embedders import OpenAIDocumentEmbedder
from haystack.components.preprocessors import EmbeddingBasedDocumentSplitter
# Create a document with content that has a clear topic shift
@@ -48,7 +48,7 @@ class EmbeddingBasedDocumentSplitter:
)
# Initialize the embedder to calculate semantic similarities
embedder = SentenceTransformersDocumentEmbedder()
embedder = OpenAIDocumentEmbedder()
# Configure the splitter with parameters that control splitting behavior
splitter = EmbeddingBasedDocumentSplitter(
haystack/components/rankers/__init__.py
-8
View File
@@ -12,8 +12,6 @@ _import_structure = {
"lost_in_the_middle": ["LostInTheMiddleRanker"],
"meta_field": ["MetaFieldRanker"],
"meta_field_grouping_ranker": ["MetaFieldGroupingRanker"],
"sentence_transformers_diversity": ["SentenceTransformersDiversityRanker"],
"sentence_transformers_similarity": ["SentenceTransformersSimilarityRanker"],
}
if TYPE_CHECKING:
@@ -21,12 +19,6 @@ if TYPE_CHECKING:
from .lost_in_the_middle import LostInTheMiddleRanker as LostInTheMiddleRanker
from .meta_field import MetaFieldRanker as MetaFieldRanker
from .meta_field_grouping_ranker import MetaFieldGroupingRanker as MetaFieldGroupingRanker
from .sentence_transformers_diversity import (
SentenceTransformersDiversityRanker as SentenceTransformersDiversityRanker,
)
from .sentence_transformers_similarity import (
SentenceTransformersSimilarityRanker as SentenceTransformersSimilarityRanker,
)
else:
sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure)
haystack/components/rankers/sentence_transformers_diversity.py
-431
View File
@@ -1,431 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from enum import Enum
from typing import Any, Literal
from haystack import Document, component, default_from_dict, default_to_dict
from haystack.lazy_imports import LazyImport
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
with LazyImport(message="Run 'pip install \"sentence-transformers>=5.0.0\"'") as torch_and_sentence_transformers_import:
import torch
from sentence_transformers import SentenceTransformer
class DiversityRankingStrategy(Enum):
"""
The strategy to use for diversity ranking.
"""
GREEDY_DIVERSITY_ORDER = "greedy_diversity_order"
MAXIMUM_MARGIN_RELEVANCE = "maximum_margin_relevance"
def __str__(self) -> str:
"""
Convert a Strategy enum to a string.
"""
return self.value
@staticmethod
def from_str(string: str) -> "DiversityRankingStrategy":
"""
Convert a string to a Strategy enum.
"""
enum_map = {e.value: e for e in DiversityRankingStrategy}
strategy = enum_map.get(string)
if strategy is None:
msg = f"Unknown strategy '{string}'. Supported strategies are: {list(enum_map.keys())}"
raise ValueError(msg)
return strategy
class DiversityRankingSimilarity(Enum):
"""
The similarity metric to use for comparing embeddings.
"""
DOT_PRODUCT = "dot_product"
COSINE = "cosine"
def __str__(self) -> str:
"""
Convert a Similarity enum to a string.
"""
return self.value
@staticmethod
def from_str(string: str) -> "DiversityRankingSimilarity":
"""
Convert a string to a Similarity enum.
"""
enum_map = {e.value: e for e in DiversityRankingSimilarity}
similarity = enum_map.get(string)
if similarity is None:
msg = f"Unknown similarity metric '{string}'. Supported metrics are: {list(enum_map.keys())}"
raise ValueError(msg)
return similarity
@component
class SentenceTransformersDiversityRanker:
"""
A Diversity Ranker based on Sentence Transformers.
Applies a document ranking algorithm based on one of the two strategies:
1. Greedy Diversity Order:
Implements a document ranking algorithm that orders documents in a way that maximizes the overall diversity
of the documents based on their similarity to the query.
It uses a pre-trained Sentence Transformers model to embed the query and
the documents.
2. Maximum Margin Relevance:
Implements a document ranking algorithm that orders documents based on their Maximum Margin Relevance (MMR)
scores.
MMR scores are calculated for each document based on their relevance to the query and diversity from already
selected documents. The algorithm iteratively selects documents based on their MMR scores, balancing between
relevance to the query and diversity from already selected documents. The 'lambda_threshold' controls the
trade-off between relevance and diversity.
Before ranking, documents are deduplicated by their id, retaining only the document with the highest score
if a score is present.
### Usage example
```python
from haystack import Document
from haystack.components.rankers import SentenceTransformersDiversityRanker
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity="cosine", strategy="greedy_diversity_order"
)
docs = [Document(content="Paris"), Document(content="Berlin")]
query = "What is the capital of germany?"
output = ranker.run(query=query, documents=docs)
docs = output["documents"]
```
"""
def __init__( # noqa: PLR0913
self,
model: str = "sentence-transformers/all-MiniLM-L6-v2",
top_k: int = 10,
device: ComponentDevice | None = None,
token: Secret | None = Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False),
similarity: str | DiversityRankingSimilarity = "cosine",
query_prefix: str = "",
query_suffix: str = "",
document_prefix: str = "",
document_suffix: str = "",
meta_fields_to_embed: list[str] | None = None,
embedding_separator: str = "\n",
strategy: str | DiversityRankingStrategy = "greedy_diversity_order",
lambda_threshold: float = 0.5,
model_kwargs: dict[str, Any] | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
config_kwargs: dict[str, Any] | None = None,
backend: Literal["torch", "onnx", "openvino"] = "torch",
) -> None:
"""
Initialize a SentenceTransformersDiversityRanker.
:param model: Local path or name of the model in Hugging Face's model hub,
such as `'sentence-transformers/all-MiniLM-L6-v2'`.
:param top_k: The maximum number of Documents to return per query.
:param device: The device on which the model is loaded. If `None`, the default device is automatically
selected.
:param token: The API token used to download private models from Hugging Face.
:param similarity: Similarity metric for comparing embeddings. Can be set to "dot_product" (default) or
"cosine".
:param query_prefix: A string to add to the beginning of the query text before ranking.
Can be used to prepend the text with an instruction, as required by some embedding models,
such as E5 and BGE.
:param query_suffix: A string to add to the end of the query text before ranking.
:param document_prefix: A string to add to the beginning of each Document text before ranking.
Can be used to prepend the text with an instruction, as required by some embedding models,
such as E5 and BGE.
:param document_suffix: A string to add to the end of each Document text before ranking.
:param meta_fields_to_embed: List of meta fields that should be embedded along with the Document content.
:param embedding_separator: Separator used to concatenate the meta fields to the Document content.
:param strategy: The strategy to use for diversity ranking. Can be either "greedy_diversity_order" or
"maximum_margin_relevance".
:param lambda_threshold: The trade-off parameter between relevance and diversity. Only used when strategy is
"maximum_margin_relevance".
:param model_kwargs:
Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
when loading the model. Refer to specific model documentation for available kwargs.
:param tokenizer_kwargs:
Additional keyword arguments for `AutoTokenizer.from_pretrained` when loading the tokenizer.
Refer to specific model documentation for available kwargs.
:param config_kwargs:
Additional keyword arguments for `AutoConfig.from_pretrained` when loading the model configuration.
:param backend:
The backend to use for the Sentence Transformers model. Choose from "torch", "onnx", or "openvino".
Refer to the [Sentence Transformers documentation](https://sbert.net/docs/sentence_transformer/usage/efficiency.html)
for more information on acceleration and quantization options.
"""
torch_and_sentence_transformers_import.check()
self.model_name_or_path = model
if top_k is None or top_k <= 0:
raise ValueError(f"top_k must be > 0, but got {top_k}")
self.top_k = top_k
self.device = ComponentDevice.resolve_device(device)
self.token = token
self.model: SentenceTransformer | None = None
self.similarity = DiversityRankingSimilarity.from_str(similarity) if isinstance(similarity, str) else similarity
self.query_prefix = query_prefix
self.document_prefix = document_prefix
self.query_suffix = query_suffix
self.document_suffix = document_suffix
self.meta_fields_to_embed = meta_fields_to_embed or []
self.embedding_separator = embedding_separator
self.strategy = DiversityRankingStrategy.from_str(strategy) if isinstance(strategy, str) else strategy
self.lambda_threshold = lambda_threshold
self._check_lambda_threshold(self.lambda_threshold, self.strategy)
self.model_kwargs = model_kwargs
self.tokenizer_kwargs = tokenizer_kwargs
self.config_kwargs = config_kwargs
self.backend = backend
def warm_up(self) -> None:
"""
Initializes the component.
"""
if self.model is None:
self.model = SentenceTransformer(
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,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
backend=self.backend,
)
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.
:returns:
Dictionary with serialized data.
"""
serialization_dict = default_to_dict(
self,
model=self.model_name_or_path,
top_k=self.top_k,
device=self.device,
token=self.token,
similarity=str(self.similarity),
query_prefix=self.query_prefix,
query_suffix=self.query_suffix,
document_prefix=self.document_prefix,
document_suffix=self.document_suffix,
meta_fields_to_embed=self.meta_fields_to_embed,
embedding_separator=self.embedding_separator,
strategy=str(self.strategy),
lambda_threshold=self.lambda_threshold,
model_kwargs=self.model_kwargs,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
backend=self.backend,
)
if serialization_dict["init_parameters"].get("model_kwargs") is not None:
serialize_hf_model_kwargs(serialization_dict["init_parameters"]["model_kwargs"])
return serialization_dict
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "SentenceTransformersDiversityRanker":
"""
Deserializes the component from a dictionary.
:param data:
The dictionary to deserialize from.
:returns:
The deserialized component.
"""
init_params = data["init_parameters"]
if init_params.get("model_kwargs") is not None:
deserialize_hf_model_kwargs(init_params["model_kwargs"])
return default_from_dict(cls, data)
def _prepare_texts_to_embed(self, documents: list[Document]) -> list[str]:
"""
Prepare the texts to embed by concatenating the Document text with the metadata fields to embed.
"""
texts_to_embed = []
for doc in documents:
meta_values_to_embed = [
str(doc.meta[key]) for key in self.meta_fields_to_embed if key in doc.meta and doc.meta[key] is not None
]
text_to_embed = (
self.document_prefix
+ self.embedding_separator.join(meta_values_to_embed + [doc.content or ""])
+ self.document_suffix
)
texts_to_embed.append(text_to_embed)
return texts_to_embed
def _greedy_diversity_order(self, query: str, documents: list[Document]) -> list[Document]:
"""
Orders the given list of documents to maximize diversity.
The algorithm first calculates embeddings for each document and the query. It starts by selecting the document
that is semantically closest to the query. Then, for each remaining document, it selects the one that, on
average, is least similar to the already selected documents. This process continues until all documents are
selected, resulting in a list where each subsequent document contributes the most to the overall diversity of
the selected set.
:param query: The search query.
:param documents: The list of Document objects to be ranked.
:return: A list of documents ordered to maximize diversity.
"""
texts_to_embed = self._prepare_texts_to_embed(documents)
doc_embeddings, query_embedding = self._embed_and_normalize(query, texts_to_embed)
n = len(documents)
selected: list[int] = []
# Compute the similarity vector between the query and documents
query_doc_sim = query_embedding @ doc_embeddings.T
# Start with the document with the highest similarity to the query
selected.append(int(torch.argmax(query_doc_sim).item()))
selected_sum = doc_embeddings[selected[0]] / n
while len(selected) < n:
# Compute mean of dot products of all selected documents and all other documents
similarities = selected_sum @ doc_embeddings.T
# Mask documents that are already selected
similarities[selected] = torch.inf
# Select the document with the lowest total similarity score
index_unselected = int(torch.argmin(similarities).item())
selected.append(index_unselected)
# It's enough just to add to the selected vectors because dot product is distributive
# It's divided by n for numerical stability
selected_sum += doc_embeddings[index_unselected] / n
ranked_docs: list[Document] = [documents[i] for i in selected]
return ranked_docs
def _embed_and_normalize(self, query: str, texts_to_embed: list[str]) -> tuple[Any, Any]:
assert self.model is not None # verified in run but mypy doesn't see it
# Calculate embeddings
doc_embeddings = self.model.encode(texts_to_embed, convert_to_tensor=True)
query_embedding = self.model.encode([self.query_prefix + query + self.query_suffix], convert_to_tensor=True)
# Normalize embeddings to unit length for computing cosine similarity
if self.similarity == DiversityRankingSimilarity.COSINE:
doc_embeddings = doc_embeddings / torch.norm(doc_embeddings, p=2, dim=-1).unsqueeze(-1)
query_embedding = query_embedding / torch.norm(query_embedding, p=2, dim=-1).unsqueeze(-1)
return doc_embeddings, query_embedding
def _maximum_margin_relevance(
self, query: str, documents: list[Document], lambda_threshold: float, top_k: int
) -> list[Document]:
"""
Orders the given list of documents according to the Maximum Margin Relevance (MMR) scores.
MMR scores are calculated for each document based on their relevance to the query and diversity from already
selected documents.
The algorithm iteratively selects documents based on their MMR scores, balancing between relevance to the query
and diversity from already selected documents. The 'lambda_threshold' controls the trade-off between relevance
and diversity.
A closer value to 0 favors diversity, while a closer value to 1 favors relevance to the query.
See : "The Use of MMR, Diversity-Based Reranking for Reordering Documents and Producing Summaries"
https://www.cs.cmu.edu/~jgc/publication/The_Use_MMR_Diversity_Based_LTMIR_1998.pdf
"""
texts_to_embed = self._prepare_texts_to_embed(documents)
doc_embeddings, query_embedding = self._embed_and_normalize(query, texts_to_embed)
top_k = min(top_k, len(documents))
selected: list[int] = []
query_similarities_as_tensor = query_embedding @ doc_embeddings.T
query_similarities = query_similarities_as_tensor.reshape(-1)
idx = int(torch.argmax(query_similarities))
selected.append(idx)
while len(selected) < top_k:
best_idx = None
best_score = -float("inf")
for idx, _ in enumerate(documents):
if idx in selected:
continue
relevance_score = query_similarities[idx]
diversity_score = max(
doc_embeddings[idx] @ doc_embeddings[j].permute(*torch.arange(doc_embeddings[j].ndim - 1, -1, -1))
for j in selected
)
mmr_score = lambda_threshold * relevance_score - (1 - lambda_threshold) * diversity_score
if mmr_score > best_score:
best_score = mmr_score
best_idx = idx
# loop condition ensures unselected docs exist with valid scores
selected.append(best_idx) # type: ignore[arg-type]
return [documents[i] for i in selected]
@staticmethod
def _check_lambda_threshold(lambda_threshold: float, strategy: DiversityRankingStrategy) -> None:
if (strategy == DiversityRankingStrategy.MAXIMUM_MARGIN_RELEVANCE) and not 0 <= lambda_threshold <= 1:
raise ValueError(f"lambda_threshold must be between 0 and 1, but got {lambda_threshold}.")
@component.output_types(documents=list[Document])
def run(
self, query: str, documents: list[Document], top_k: int | None = None, lambda_threshold: float | None = None
) -> dict[str, list[Document]]:
"""
Rank the documents based on their diversity.
:param query: The search query.
:param documents: List of Document objects to be ranker.
:param top_k: Optional. An integer to override the top_k set during initialization.
:param lambda_threshold: Override the trade-off parameter between relevance and diversity. Only used when
strategy is "maximum_margin_relevance".
:returns: A dictionary with the following key:
- `documents`: List of Document objects that have been selected based on the diversity ranking.
:raises ValueError: If the top_k value is less than or equal to 0.
"""
if self.model is None:
self.warm_up()
if not documents:
return {"documents": []}
if top_k is None:
top_k = self.top_k
if top_k <= 0:
raise ValueError(f"top_k must be > 0, but got {top_k}")
deduplicated_documents = _deduplicate_documents(documents)
if self.strategy == DiversityRankingStrategy.MAXIMUM_MARGIN_RELEVANCE:
if lambda_threshold is None:
lambda_threshold = self.lambda_threshold
self._check_lambda_threshold(lambda_threshold, self.strategy)
re_ranked_docs = self._maximum_margin_relevance(
query=query, documents=deduplicated_documents, lambda_threshold=lambda_threshold, top_k=top_k
)
else:
re_ranked_docs = self._greedy_diversity_order(query=query, documents=deduplicated_documents)
return {"documents": re_ranked_docs[:top_k]}
haystack/components/rankers/sentence_transformers_similarity.py
-296
View File
@@ -1,296 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from dataclasses import replace
from pathlib import Path
from typing import Any, Literal
from haystack import Document, component, default_from_dict, default_to_dict
from haystack.lazy_imports import LazyImport
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
with LazyImport(message="Run 'pip install \"sentence-transformers>=5.0.0\"'") as torch_and_sentence_transformers_import:
from sentence_transformers import CrossEncoder
from torch.nn import Identity, Sigmoid
@component
class SentenceTransformersSimilarityRanker:
"""
Ranks documents based on their semantic similarity to the query.
It uses a pre-trained cross-encoder model from Hugging Face to embed the query and the documents.
### Usage example
```python
from haystack import Document
from haystack.components.rankers import SentenceTransformersSimilarityRanker
ranker = SentenceTransformersSimilarityRanker()
docs = [Document(content="Paris"), Document(content="Berlin")]
query = "City in Germany"
result = ranker.run(query=query, documents=docs)
docs = result["documents"]
print(docs[0].content)
```
"""
def __init__( # noqa: PLR0913
self,
*,
model: str | Path = "cross-encoder/ms-marco-MiniLM-L-6-v2",
device: ComponentDevice | None = None,
token: Secret | None = Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False),
top_k: int = 10,
query_prefix: str = "",
query_suffix: str = "",
document_prefix: str = "",
document_suffix: str = "",
meta_fields_to_embed: list[str] | None = None,
embedding_separator: str = "\n",
scale_score: bool = True,
score_threshold: float | None = None,
trust_remote_code: bool = False,
model_kwargs: dict[str, Any] | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
config_kwargs: dict[str, Any] | None = None,
backend: Literal["torch", "onnx", "openvino"] = "torch",
batch_size: int = 16,
) -> None:
"""
Creates an instance of SentenceTransformersSimilarityRanker.
:param model:
The ranking model. Pass a local path or the Hugging Face model name of a cross-encoder model.
:param device:
The device on which the model is loaded. If `None`, the default device is automatically selected.
:param token:
The API token to download private models from Hugging Face.
:param top_k:
The maximum number of documents to return per query.
:param query_prefix:
A string to add at the beginning of the query text before ranking.
Use it to prepend the text with an instruction, as required by reranking models like `bge`.
:param query_suffix:
A string to add at the end of the query text before ranking.
Use it to append the text with an instruction, as required by reranking models like `qwen`.
:param document_prefix:
A string to add at the beginning of each document before ranking. You can use it to prepend the document
with an instruction, as required by embedding models like `bge`.
:param document_suffix:
A string to add at the end of each document before ranking. You can use it to append the document
with an instruction, as required by embedding models like `qwen`.
:param meta_fields_to_embed:
List of metadata fields to embed with the document.
:param embedding_separator:
Separator to concatenate metadata fields to the document.
:param scale_score:
If `True`, scales the raw logit predictions using a Sigmoid activation function.
If `False`, disables scaling of the raw logit predictions.
:param score_threshold:
Use it to return documents with a score above this threshold only.
:param trust_remote_code:
If `False`, allows only Hugging Face verified model architectures.
If `True`, allows custom models and scripts.
:param model_kwargs:
Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
when loading the model. Refer to specific model documentation for available kwargs.
:param tokenizer_kwargs:
Additional keyword arguments for `AutoTokenizer.from_pretrained` when loading the tokenizer.
Refer to specific model documentation for available kwargs.
:param config_kwargs:
Additional keyword arguments for `AutoConfig.from_pretrained` when loading the model configuration.
:param backend:
The backend to use for the Sentence Transformers model. Choose from "torch", "onnx", or "openvino".
Refer to the [Sentence Transformers documentation](https://sbert.net/docs/sentence_transformer/usage/efficiency.html)
for more information on acceleration and quantization options.
:param batch_size:
The batch size to use for inference. The higher the batch size, the more memory is required.
If you run into memory issues, reduce the batch size.
:raises ValueError:
If `top_k` is not > 0.
"""
torch_and_sentence_transformers_import.check()
if top_k <= 0:
raise ValueError(f"top_k must be > 0, but got {top_k}")
self.model = str(model)
self._cross_encoder = None
self.query_prefix = query_prefix
self.query_suffix = query_suffix
self.document_prefix = document_prefix
self.document_suffix = document_suffix
self.device = ComponentDevice.resolve_device(device)
self.top_k = top_k
self.token = token
self.meta_fields_to_embed = meta_fields_to_embed or []
self.embedding_separator = embedding_separator
self.scale_score = scale_score
self.score_threshold = score_threshold
self.trust_remote_code = trust_remote_code
self.model_kwargs = model_kwargs
self.tokenizer_kwargs = tokenizer_kwargs
self.config_kwargs = config_kwargs
self.backend = backend
self.batch_size = batch_size
def _get_telemetry_data(self) -> dict[str, Any]:
"""
Data that is sent to Posthog for usage analytics.
"""
return {"model": self.model}
def warm_up(self) -> None:
"""
Initializes the component.
"""
if self._cross_encoder is None:
self._cross_encoder = CrossEncoder(
model_name_or_path=self.model,
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,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
backend=self.backend,
)
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.
:returns:
Dictionary with serialized data.
"""
serialization_dict = default_to_dict(
self,
device=self.device,
model=self.model,
token=self.token,
top_k=self.top_k,
query_prefix=self.query_prefix,
query_suffix=self.query_suffix,
document_prefix=self.document_prefix,
document_suffix=self.document_suffix,
meta_fields_to_embed=self.meta_fields_to_embed,
embedding_separator=self.embedding_separator,
scale_score=self.scale_score,
score_threshold=self.score_threshold,
trust_remote_code=self.trust_remote_code,
model_kwargs=self.model_kwargs,
tokenizer_kwargs=self.tokenizer_kwargs,
config_kwargs=self.config_kwargs,
backend=self.backend,
batch_size=self.batch_size,
)
if serialization_dict["init_parameters"].get("model_kwargs") is not None:
serialize_hf_model_kwargs(serialization_dict["init_parameters"]["model_kwargs"])
return serialization_dict
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "SentenceTransformersSimilarityRanker":
"""
Deserializes the component from a dictionary.
:param data:
Dictionary to deserialize from.
:returns:
Deserialized component.
"""
init_params = data["init_parameters"]
if init_params.get("model_kwargs") is not None:
deserialize_hf_model_kwargs(init_params["model_kwargs"])
return default_from_dict(cls, data)
@component.output_types(documents=list[Document])
def run(
self,
*,
query: str,
documents: list[Document],
top_k: int | None = None,
scale_score: bool | None = None,
score_threshold: float | None = None,
) -> dict[str, list[Document]]:
"""
Returns a list of documents ranked by their similarity to the given query.
Before ranking, documents are deduplicated by their id, retaining only the document with the highest score
if a score is present.
:param query:
The input query to compare the documents to.
:param documents:
A list of documents to be ranked.
:param top_k:
The maximum number of documents to return.
:param scale_score:
If `True`, scales the raw logit predictions using a Sigmoid activation function.
If `False`, disables scaling of the raw logit predictions.
If set, overrides the value set at initialization.
:param score_threshold:
Use it to return documents only with a score above this threshold.
If set, overrides the value set at initialization.
:returns:
A dictionary with the following keys:
- `documents`: A list of documents closest to the query, sorted from most similar to least similar.
:raises ValueError:
If `top_k` is not > 0.
"""
if self._cross_encoder is None:
self.warm_up()
if not documents:
return {"documents": []}
top_k = top_k or self.top_k
scale_score = scale_score or self.scale_score
score_threshold = score_threshold or self.score_threshold
if top_k <= 0:
raise ValueError(f"top_k must be > 0, but got {top_k}")
deduplicated_documents = _deduplicate_documents(documents)
prepared_query = self.query_prefix + query + self.query_suffix
prepared_documents = []
for doc in deduplicated_documents:
meta_values_to_embed = [
str(doc.meta[key]) for key in self.meta_fields_to_embed if key in doc.meta and doc.meta[key] is not None
]
prepared_documents.append(
self.document_prefix
+ self.embedding_separator.join(meta_values_to_embed + [doc.content or ""])
+ self.document_suffix
)
activation_fn = Sigmoid() if scale_score else Identity()
# mypy doesn't know this is set in warm_up
ranking_result = self._cross_encoder.rank( # type: ignore[attr-defined]
query=prepared_query,
documents=prepared_documents,
batch_size=self.batch_size,
activation_fn=activation_fn,
convert_to_numpy=True,
return_documents=False,
)
ranked_docs = []
for el in ranking_result:
index = el["corpus_id"]
score = float(el["score"])
ranked_docs.append(replace(deduplicated_documents[index], score=score))
if score_threshold is not None:
ranked_docs = [doc for doc in ranked_docs if doc.score >= score_threshold]
return {"documents": ranked_docs[:top_k]}
haystack/components/retrievers/in_memory/embedding_retriever.py
+3 -3
View File
@@ -23,7 +23,7 @@ class InMemoryEmbeddingRetriever:
### Usage example
```python
from haystack import Document
from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
from haystack.components.embedders import OpenAIDocumentEmbedder, OpenAITextEmbedder
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from haystack.document_stores.in_memory import InMemoryDocumentStore
@@ -31,7 +31,7 @@ class InMemoryEmbeddingRetriever:
Document(content="Python is a popular programming language"),
Document(content="python ist eine beliebte Programmiersprache"),
]
doc_embedder = SentenceTransformersDocumentEmbedder()
doc_embedder = OpenAIDocumentEmbedder()
docs_with_embeddings = doc_embedder.run(docs)["documents"]
doc_store = InMemoryDocumentStore()
@@ -39,7 +39,7 @@ class InMemoryEmbeddingRetriever:
retriever = InMemoryEmbeddingRetriever(doc_store)
query="Programmiersprache"
text_embedder = SentenceTransformersTextEmbedder()
text_embedder = OpenAITextEmbedder()
query_embedding = text_embedder.run(query)["embedding"]
result = retriever.run(query_embedding=query_embedding)
haystack/components/retrievers/multi_query_embedding_retriever.py
+4 -4
View File
@@ -29,8 +29,8 @@ class MultiQueryEmbeddingRetriever:
from haystack import Document
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.document_stores.types import DuplicatePolicy
from haystack.components.embedders import SentenceTransformersTextEmbedder
from haystack.components.embedders import SentenceTransformersDocumentEmbedder
from haystack.components.embedders import OpenAITextEmbedder
from haystack.components.embedders import OpenAIDocumentEmbedder
from haystack.components.retrievers import InMemoryEmbeddingRetriever
from haystack.components.writers import DocumentWriter
from haystack.components.retrievers import MultiQueryEmbeddingRetriever
@@ -46,14 +46,14 @@ class MultiQueryEmbeddingRetriever:
# Populate the document store
doc_store = InMemoryDocumentStore()
doc_embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
doc_embedder = OpenAIDocumentEmbedder()
doc_writer = DocumentWriter(document_store=doc_store, policy=DuplicatePolicy.SKIP)
documents = doc_embedder.run(documents)["documents"]
doc_writer.run(documents=documents)
# Run the multi-query retriever
in_memory_retriever = InMemoryEmbeddingRetriever(document_store=doc_store, top_k=1)
query_embedder = SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
query_embedder = OpenAITextEmbedder()
multi_query_retriever = MultiQueryEmbeddingRetriever(
retriever=in_memory_retriever,
haystack/components/retrievers/multi_retriever.py
+3 -3
View File
@@ -39,7 +39,7 @@ class MultiRetriever:
from haystack.document_stores.types import DuplicatePolicy
from haystack.components.retrievers import InMemoryBM25Retriever, InMemoryEmbeddingRetriever
from haystack.components.retrievers import TextEmbeddingRetriever, MultiRetriever
from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder
from haystack.components.embedders import OpenAITextEmbedder, OpenAIDocumentEmbedder
from haystack.components.writers import DocumentWriter
documents = [
@@ -50,7 +50,7 @@ class MultiRetriever:
# Populate the document store
doc_store = InMemoryDocumentStore()
doc_embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
doc_embedder = OpenAIDocumentEmbedder()
doc_writer = DocumentWriter(document_store=doc_store, policy=DuplicatePolicy.SKIP)
doc_writer.run(documents=doc_embedder.run(documents)["documents"])
@@ -60,7 +60,7 @@ class MultiRetriever:
"bm25": InMemoryBM25Retriever(document_store=doc_store),
"embedding": TextEmbeddingRetriever(
retriever=InMemoryEmbeddingRetriever(document_store=doc_store),
text_embedder=SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"),
text_embedder=OpenAITextEmbedder(),
),
},
top_k=3,
haystack/components/retrievers/text_embedding_retriever.py
+3 -3
View File
@@ -26,7 +26,7 @@ class TextEmbeddingRetriever:
from haystack import Document
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.document_stores.types import DuplicatePolicy
from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder
from haystack.components.embedders import OpenAITextEmbedder, OpenAIDocumentEmbedder
from haystack.components.retrievers import InMemoryEmbeddingRetriever, TextEmbeddingRetriever
from haystack.components.writers import DocumentWriter
@@ -41,14 +41,14 @@ class TextEmbeddingRetriever:
# Populate the document store
doc_store = InMemoryDocumentStore()
doc_embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
doc_embedder = OpenAIDocumentEmbedder()
doc_writer = DocumentWriter(document_store=doc_store, policy=DuplicatePolicy.SKIP)
documents = doc_embedder.run(documents)["documents"]
doc_writer.run(documents=documents)
# Run the retriever
in_memory_retriever = InMemoryEmbeddingRetriever(document_store=doc_store, top_k=1)
text_embedder = SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
text_embedder = OpenAITextEmbedder()
retriever = TextEmbeddingRetriever(retriever=in_memory_retriever, text_embedder=text_embedder)
result = retriever.run(query="Geothermal energy")
haystack/tools/pipeline_tool.py
+3 -8
View File
@@ -38,10 +38,7 @@ class PipelineTool(ComponentTool):
from haystack import Document, Pipeline
from haystack.dataclasses import ChatMessage
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.components.embedders.sentence_transformers_text_embedder import SentenceTransformersTextEmbedder
from haystack.components.embedders.sentence_transformers_document_embedder import (
SentenceTransformersDocumentEmbedder
)
from haystack.components.embedders import OpenAITextEmbedder, OpenAIDocumentEmbedder
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.retrievers import InMemoryEmbeddingRetriever
from haystack.components.agents import Agent
@@ -49,7 +46,7 @@ class PipelineTool(ComponentTool):
# Initialize a document store and add some documents
document_store = InMemoryDocumentStore()
document_embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
document_embedder = OpenAIDocumentEmbedder()
documents = [
Document(content="Nikola Tesla was a Serbian-American inventor and electrical engineer."),
Document(
@@ -62,9 +59,7 @@ class PipelineTool(ComponentTool):
# Build a simple retrieval pipeline
retrieval_pipeline = Pipeline()
retrieval_pipeline.add_component(
"embedder", SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
)
retrieval_pipeline.add_component("embedder", OpenAITextEmbedder())
retrieval_pipeline.add_component("retriever", InMemoryEmbeddingRetriever(document_store=document_store))
retrieval_pipeline.connect("embedder.embedding", "retriever.query_embedding")
haystack/utils/hf.py
-49
View File
@@ -7,15 +7,10 @@ from typing import Any
from haystack import logging
from haystack.dataclasses import ChatMessage, ImageContent, ReasoningContent, TextContent
from haystack.lazy_imports import LazyImport
from haystack.utils.auth import Secret
from haystack.utils.device import ComponentDevice
with LazyImport(message="Run 'pip install \"transformers[torch]\"'") as torch_import:
import torch
with LazyImport(message="Run 'pip install \"huggingface_hub>=0.27.0\"'") as huggingface_hub_import:
from huggingface_hub import model_info
logger = logging.getLogger(__name__)
@@ -56,50 +51,6 @@ def deserialize_hf_model_kwargs(kwargs: dict[str, Any]) -> None:
deserialize_hf_model_kwargs(v)
def resolve_hf_pipeline_kwargs(
huggingface_pipeline_kwargs: dict[str, Any],
model: str,
task: str | None,
supported_tasks: list[str],
device: ComponentDevice | None,
token: Secret | None,
) -> dict[str, Any]:
"""
Resolve the HuggingFace pipeline keyword arguments based on explicit user inputs.
:param huggingface_pipeline_kwargs: Dictionary containing keyword arguments used to initialize a
Hugging Face pipeline.
:param model: The name or path of a Hugging Face model for on the HuggingFace Hub.
:param task: The task for the Hugging Face pipeline.
:param supported_tasks: The list of supported tasks to check the task of the model against. If the task of the model
is not present within this list then a ValueError is thrown.
:param device: The device on which the model is loaded. If `None`, the default device is automatically
selected. If a device/device map is specified in `huggingface_pipeline_kwargs`, it overrides this parameter.
:param token: The token to use as HTTP bearer authorization for remote files.
If the token is also specified in the `huggingface_pipeline_kwargs`, this parameter will be ignored.
"""
huggingface_hub_import.check()
token = token.resolve_value() if token else None
# check if the huggingface_pipeline_kwargs contain the essential parameters
# otherwise, populate them with values from other init parameters
huggingface_pipeline_kwargs.setdefault("model", model)
huggingface_pipeline_kwargs.setdefault("token", token)
device = ComponentDevice.resolve_device(device)
device.update_hf_kwargs(huggingface_pipeline_kwargs, overwrite=False)
# task identification and validation
task = task or huggingface_pipeline_kwargs.get("task")
if task is None and isinstance(huggingface_pipeline_kwargs["model"], str):
task = model_info(huggingface_pipeline_kwargs["model"], token=huggingface_pipeline_kwargs["token"]).pipeline_tag
if task not in supported_tasks:
raise ValueError(f"Task '{task}' is not supported. The supported tasks are: {', '.join(supported_tasks)}.")
huggingface_pipeline_kwargs["task"] = task
return huggingface_pipeline_kwargs
def convert_message_to_hf_format(message: ChatMessage) -> dict[str, Any]:
"""
Convert a message to the format expected by Hugging Face.
pydoc/embedders_api.yml
-5
View File
@@ -6,11 +6,6 @@ loaders:
"azure_text_embedder",
"openai_document_embedder",
"openai_text_embedder",
"sentence_transformers_document_embedder",
"sentence_transformers_text_embedder",
"sentence_transformers_sparse_document_embedder",
"sentence_transformers_sparse_text_embedder",
"image/sentence_transformers_doc_image_embedder",
]
processors:
- type: filter
pydoc/rankers_api.yml
+1 -2
View File
@@ -1,7 +1,6 @@
loaders:
- search_path: [../haystack/components/rankers]
modules: ["llm_ranker", "lost_in_the_middle", "meta_field", "meta_field_grouping_ranker",
"sentence_transformers_diversity", "sentence_transformers_similarity"]
modules: ["llm_ranker", "lost_in_the_middle", "meta_field", "meta_field_grouping_ranker"]
processors:
- type: filter
documented_only: true
pyproject.toml
+2 -2
View File
@@ -92,8 +92,8 @@ dependencies = [
"pandas", # AzureOCRDocumentConverter, CSVDocumentCleaner, CSVDocumentSplitter,
# EvaluationRunResult, XLSXToDocument, and pipeline tests
"transformers[torch, sentencepiece]>=4.57", # NamedEntityExtractor, SASEvaluator, and TopPSampler
"sentence-transformers>=5.0.0", # Sentence Transformers Embedders, Rankers, and SASEvaluator
"transformers[torch, sentencepiece]>=4.57", # SASEvaluator and TopPSampler
"sentence-transformers>=5.0.0", # SASEvaluator
"langdetect", # TextLanguageRouter and DocumentLanguageClassifier
"openai-whisper>=20231106", # LocalWhisperTranscriber
"arrow>=1.3.0", # Jinja2TimeExtension
releasenotes/notes/remove-sentence-transformers-components-3f8fa21ffa1aede6.yaml
+39
View File
@@ -0,0 +1,39 @@
---
upgrade:
- |
The Sentence Transformers components have been moved out of Haystack into the
``sentence-transformers-haystack`` integration package. This includes
``SentenceTransformersTextEmbedder``, ``SentenceTransformersDocumentEmbedder``,
``SentenceTransformersSparseTextEmbedder``, ``SentenceTransformersSparseDocumentEmbedder``,
``SentenceTransformersDocumentImageEmbedder``, ``SentenceTransformersSimilarityRanker``, and
``SentenceTransformersDiversityRanker``. Install the new package with
``pip install sentence-transformers-haystack`` and update your imports.
Before:
.. code:: python
from haystack.components.embedders import (
SentenceTransformersTextEmbedder,
SentenceTransformersDocumentEmbedder,
SentenceTransformersSparseTextEmbedder,
SentenceTransformersSparseDocumentEmbedder,
)
from haystack.components.embedders.image import SentenceTransformersDocumentImageEmbedder
from haystack.components.rankers import SentenceTransformersSimilarityRanker, SentenceTransformersDiversityRanker
After:
.. code:: python
from haystack_integrations.components.embedders.sentence_transformers import (
SentenceTransformersTextEmbedder,
SentenceTransformersDocumentEmbedder,
SentenceTransformersSparseTextEmbedder,
SentenceTransformersSparseDocumentEmbedder,
SentenceTransformersDocumentImageEmbedder,
)
from haystack_integrations.components.rankers.sentence_transformers import (
SentenceTransformersSimilarityRanker,
SentenceTransformersDiversityRanker,
)
test/components/embedders/image/test_sentence_transformers_doc_image_embedder.py
-345
View File
@@ -1,345 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
import glob
import random
import sys
from unittest.mock import MagicMock, patch
import pytest
import torch
from PIL import Image
from haystack import Document
from haystack.components.embedders.image.sentence_transformers_doc_image_embedder import (
SentenceTransformersDocumentImageEmbedder,
)
from haystack.utils.auth import Secret
from haystack.utils.device import ComponentDevice
IMPORT_PATH = "haystack.components.embedders.image.sentence_transformers_doc_image_embedder"
class TestSentenceTransformersDocumentImageEmbedder:
def test_init_default(self):
embedder = SentenceTransformersDocumentImageEmbedder(model="model")
assert embedder.file_path_meta_field == "file_path"
assert embedder.root_path == ""
assert embedder.model == "model"
assert embedder.device == ComponentDevice.resolve_device(None)
assert embedder.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert embedder.batch_size == 32
assert embedder.progress_bar is True
assert embedder.trust_remote_code is False
assert embedder.local_files_only is False
assert embedder.precision == "float32"
assert embedder._embedding_backend is None
def test_init_with_parameters(self):
embedder = SentenceTransformersDocumentImageEmbedder(
file_path_meta_field="custom_file_path",
root_path="root_path",
model="model",
device=ComponentDevice.from_str("cuda:0"),
token=Secret.from_token("fake-api-token"),
batch_size=64,
progress_bar=False,
trust_remote_code=True,
local_files_only=True,
precision="int8",
)
assert embedder.file_path_meta_field == "custom_file_path"
assert embedder.root_path == "root_path"
assert embedder.model == "model"
assert embedder.device == ComponentDevice.from_str("cuda:0")
assert embedder.token == Secret.from_token("fake-api-token")
assert embedder.batch_size == 64
assert embedder.progress_bar is False
assert embedder.trust_remote_code
assert embedder.local_files_only
assert embedder.precision == "int8"
assert embedder.backend == "torch"
assert embedder.model_kwargs is None
assert embedder.tokenizer_kwargs is None
assert embedder.config_kwargs is None
assert embedder.encode_kwargs is None
assert embedder._embedding_backend is None
def test_to_dict(self):
component = SentenceTransformersDocumentImageEmbedder(
model="model", device=ComponentDevice.from_str("cpu"), model_kwargs={"torch_dtype": "torch.float32"}
)
data = component.to_dict()
assert data == {
"type": f"{IMPORT_PATH}.SentenceTransformersDocumentImageEmbedder",
"init_parameters": {
"file_path_meta_field": "file_path",
"root_path": "",
"model": "model",
"device": ComponentDevice.from_str("cpu").to_dict(),
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"batch_size": 32,
"progress_bar": True,
"normalize_embeddings": False,
"trust_remote_code": False,
"local_files_only": False,
"model_kwargs": {"torch_dtype": "torch.float32"},
"tokenizer_kwargs": None,
"encode_kwargs": None,
"config_kwargs": None,
"precision": "float32",
"backend": "torch",
},
}
def test_from_dict(self):
init_parameters = {
"file_path_meta_field": "custom_file_path",
"root_path": "root_path",
"model": "model",
"device": ComponentDevice.from_str("cuda:0").to_dict(),
"token": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
"batch_size": 64,
"progress_bar": False,
"normalize_embeddings": True,
"trust_remote_code": True,
"local_files_only": True,
"model_kwargs": {"torch_dtype": "torch.float32"},
"tokenizer_kwargs": {"model_max_length": 512},
"config_kwargs": {"use_memory_efficient_attention": True},
"precision": "int8",
}
component = SentenceTransformersDocumentImageEmbedder.from_dict(
{"type": f"{IMPORT_PATH}.SentenceTransformersDocumentImageEmbedder", "init_parameters": init_parameters}
)
assert component.file_path_meta_field == "custom_file_path"
assert component.root_path == "root_path"
assert component.model == "model"
assert component.device == ComponentDevice.from_str("cuda:0")
assert component.token == Secret.from_env_var("ENV_VAR", strict=False)
assert component.batch_size == 64
assert component.progress_bar is False
assert component.normalize_embeddings is True
assert component.trust_remote_code
assert component.local_files_only
assert component.model_kwargs == {"torch_dtype": torch.float32}
assert component.tokenizer_kwargs == {"model_max_length": 512}
assert component.config_kwargs == {"use_memory_efficient_attention": True}
assert component.precision == "int8"
def test_from_dict_none_device(self):
init_parameters = {
"file_path_meta_field": "custom_file_path",
"root_path": "root_path",
"model": "model",
"device": None,
"token": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
"batch_size": 64,
"progress_bar": False,
"normalize_embeddings": False,
"trust_remote_code": True,
"local_files_only": False,
"precision": "float32",
}
component = SentenceTransformersDocumentImageEmbedder.from_dict(
{"type": f"{IMPORT_PATH}.SentenceTransformersDocumentImageEmbedder", "init_parameters": init_parameters}
)
assert component.file_path_meta_field == "custom_file_path"
assert component.root_path == "root_path"
assert component.model == "model"
assert component.device == ComponentDevice.resolve_device(None)
assert component.token == Secret.from_env_var("ENV_VAR", strict=False)
assert component.batch_size == 64
assert component.progress_bar is False
assert component.trust_remote_code
assert component.local_files_only is False
assert component.precision == "float32"
@patch(f"{IMPORT_PATH}._SentenceTransformersEmbeddingBackendFactory")
def test_warmup(self, mocked_factory):
embedder = SentenceTransformersDocumentImageEmbedder(
model="model",
token=None,
device=ComponentDevice.from_str("cpu"),
tokenizer_kwargs={"model_max_length": 512},
config_kwargs={"use_memory_efficient_attention": True},
)
mocked_factory.get_embedding_backend.assert_not_called()
embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="model",
device="cpu",
auth_token=None,
trust_remote_code=False,
local_files_only=False,
model_kwargs=None,
tokenizer_kwargs={"model_max_length": 512},
config_kwargs={"use_memory_efficient_attention": True},
backend="torch",
)
@patch(f"{IMPORT_PATH}._SentenceTransformersEmbeddingBackendFactory")
def test_warmup_doesnt_reload(self, mocked_factory):
embedder = SentenceTransformersDocumentImageEmbedder(model="model")
mocked_factory.get_embedding_backend.assert_not_called()
embedder.warm_up()
embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once()
def test_run(self, test_files_path):
embedder = SentenceTransformersDocumentImageEmbedder(model="model")
embedder._embedding_backend = MagicMock()
embedder._embedding_backend.embed = lambda data, **_: [
[random.random() for _ in range(16)] for _ in range(len(data))
]
image_paths = glob.glob(str(test_files_path / "images" / "*.*")) + glob.glob(
str(test_files_path / "pdf" / "*.pdf")
)
documents = []
for i, path in enumerate(image_paths):
document = Document(content=f"document number {i}", meta={"file_path": path})
if path.endswith(".pdf"):
document.meta["page_number"] = 1
documents.append(document)
result = embedder.run(documents=documents)
assert isinstance(result["documents"], list)
assert len(result["documents"]) == len(documents)
for doc, new_doc in zip(documents, result["documents"], strict=True):
assert doc.embedding is None
assert new_doc is not doc
assert isinstance(new_doc, Document)
assert isinstance(new_doc.embedding, list)
assert isinstance(new_doc.embedding[0], float)
assert "embedding_source" not in doc.meta
assert "embedding_source" in new_doc.meta
assert new_doc.meta["embedding_source"]["type"] == "image"
assert "file_path_meta_field" in new_doc.meta["embedding_source"]
def test_run_wrong_input_format(self):
embedder = SentenceTransformersDocumentImageEmbedder(model="model")
string_input = "text"
list_integers_input = [1, 2, 3]
with pytest.raises(
TypeError, match="SentenceTransformersDocumentImageEmbedder expects a list of Documents as input"
):
embedder.run(documents=string_input)
with pytest.raises(
TypeError, match="SentenceTransformersDocumentImageEmbedder expects a list of Documents as input"
):
embedder.run(documents=list_integers_input)
@patch(f"{IMPORT_PATH}._SentenceTransformersEmbeddingBackendFactory")
def test_model_onnx_backend(self, mocked_factory):
onnx_embedder = SentenceTransformersDocumentImageEmbedder(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cpu"),
model_kwargs={"file_name": "onnx/model.onnx"},
# setting the path isn't necessary if the repo contains a "onnx/model.onnx" file
# but this is to prevent a HF warning
backend="onnx",
)
onnx_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="sentence-transformers/all-MiniLM-L6-v2",
device="cpu",
auth_token=None,
trust_remote_code=False,
local_files_only=False,
model_kwargs={"file_name": "onnx/model.onnx"},
tokenizer_kwargs=None,
config_kwargs=None,
backend="onnx",
)
@patch(f"{IMPORT_PATH}._SentenceTransformersEmbeddingBackendFactory")
def test_model_openvino_backend(self, mocked_factory):
openvino_embedder = SentenceTransformersDocumentImageEmbedder(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cpu"),
model_kwargs={"file_name": "openvino/openvino_model.xml"},
# setting the path isn't necessary if the repo contains a "openvino/openvino_model.xml" file
# but this is to prevent a HF warning
backend="openvino",
)
openvino_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="sentence-transformers/all-MiniLM-L6-v2",
device="cpu",
auth_token=None,
trust_remote_code=False,
local_files_only=False,
model_kwargs={"file_name": "openvino/openvino_model.xml"},
tokenizer_kwargs=None,
config_kwargs=None,
backend="openvino",
)
@patch(f"{IMPORT_PATH}._extract_image_sources_info")
@patch(f"{IMPORT_PATH}._batch_convert_pdf_pages_to_images")
@patch("PIL.Image.open")
def test_run_none_images(
self, mocked_pil_open, mocked_batch_convert_pdf_pages_to_images, mocked_extract_image_sources_info
):
embedder = SentenceTransformersDocumentImageEmbedder(model="model")
embedder._embedding_backend = MagicMock()
mocked_extract_image_sources_info.return_value = [
{"path": "doc1.pdf", "mime_type": "application/pdf", "page_number": 999}, # Page 999 doesn't exist
{"path": "image1.jpg", "mime_type": "image/jpeg"},
]
mocked_batch_convert_pdf_pages_to_images.return_value = {} # Empty dict because page was skipped
mocked_pil_open.return_value = Image.new("RGB", (100, 100))
documents = [
Document(content="PDF 1", meta={"file_path": "doc1.pdf", "page_number": 999}),
Document(content="Image 1", meta={"file_path": "image1.jpg"}),
]
with pytest.raises(RuntimeError, match="Conversion failed for some documents."):
embedder.run(documents=documents)
@pytest.mark.integration
@pytest.mark.slow
@pytest.mark.skipif(
sys.platform == "darwin",
reason=(
"This model does not play well with GitHub macOS runners and"
"we prefer to avoid altering PYTORCH_MPS_HIGH_WATERMARK_RATIO"
),
)
def test_live_run(self, test_files_path, del_hf_env_vars):
embedder = SentenceTransformersDocumentImageEmbedder(model="sentence-transformers/clip-ViT-B-32")
documents = [
Document(
content="PDF document",
meta={"file_path": str(test_files_path / "pdf" / "sample_pdf_1.pdf"), "page_number": 1},
),
Document(content="Image document", meta={"file_path": str(test_files_path / "images" / "apple.jpg")}),
]
result = embedder.run(documents=documents)
assert len(result["documents"]) == len(documents)
for doc, new_doc in zip(documents, result["documents"], strict=True):
assert doc.embedding is None
assert new_doc is not doc
assert isinstance(new_doc, Document)
assert isinstance(new_doc.embedding, list)
assert len(new_doc.embedding) == 512
assert all(isinstance(x, float) for x in new_doc.embedding)
assert "embedding_source" not in doc.meta
assert "embedding_source" in new_doc.meta
assert new_doc.meta["embedding_source"]["type"] == "image"
assert "file_path_meta_field" in new_doc.meta["embedding_source"]
test/components/embedders/test_sentence_transformers_document_embedder.py
-493
View File
@@ -1,493 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
import random
from unittest.mock import MagicMock, patch
import pytest
import torch
from haystack import Document
from haystack.components.embedders.sentence_transformers_document_embedder import SentenceTransformersDocumentEmbedder
from haystack.utils import ComponentDevice, Secret
class TestSentenceTransformersDocumentEmbedder:
def test_init_default(self):
embedder = SentenceTransformersDocumentEmbedder(model="model")
assert embedder.model == "model"
assert embedder.device == ComponentDevice.resolve_device(None)
assert embedder.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert embedder.prefix == ""
assert embedder.suffix == ""
assert embedder.batch_size == 32
assert embedder.progress_bar is True
assert embedder.normalize_embeddings is False
assert embedder.meta_fields_to_embed == []
assert embedder.embedding_separator == "\n"
assert embedder.trust_remote_code is False
assert embedder.revision is None
assert embedder.local_files_only is False
assert embedder.truncate_dim is None
assert embedder.precision == "float32"
def test_init_with_parameters(self):
embedder = SentenceTransformersDocumentEmbedder(
model="model",
device=ComponentDevice.from_str("cuda:0"),
token=Secret.from_token("fake-api-token"),
prefix="prefix",
suffix="suffix",
batch_size=64,
progress_bar=False,
normalize_embeddings=True,
meta_fields_to_embed=["test_field"],
embedding_separator=" | ",
trust_remote_code=True,
revision="v1.0",
local_files_only=True,
truncate_dim=256,
precision="int8",
)
assert embedder.model == "model"
assert embedder.device == ComponentDevice.from_str("cuda:0")
assert embedder.token == Secret.from_token("fake-api-token")
assert embedder.prefix == "prefix"
assert embedder.suffix == "suffix"
assert embedder.batch_size == 64
assert embedder.progress_bar is False
assert embedder.normalize_embeddings is True
assert embedder.meta_fields_to_embed == ["test_field"]
assert embedder.embedding_separator == " | "
assert embedder.trust_remote_code
assert embedder.revision == "v1.0"
assert embedder.local_files_only
assert embedder.truncate_dim == 256
assert embedder.precision == "int8"
def test_to_dict(self):
component = SentenceTransformersDocumentEmbedder(model="model", device=ComponentDevice.from_str("cpu"))
data = component.to_dict()
assert data == {
"type": "haystack.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", # noqa: E501
"init_parameters": {
"model": "model",
"device": ComponentDevice.from_str("cpu").to_dict(),
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"prefix": "",
"suffix": "",
"batch_size": 32,
"progress_bar": True,
"normalize_embeddings": False,
"embedding_separator": "\n",
"meta_fields_to_embed": [],
"trust_remote_code": False,
"revision": None,
"local_files_only": False,
"truncate_dim": None,
"model_kwargs": None,
"tokenizer_kwargs": None,
"encode_kwargs": None,
"config_kwargs": None,
"precision": "float32",
"backend": "torch",
},
}
def test_to_dict_with_custom_init_parameters(self):
component = SentenceTransformersDocumentEmbedder(
model="model",
device=ComponentDevice.from_str("cuda:0"),
token=Secret.from_env_var("ENV_VAR", strict=False),
prefix="prefix",
suffix="suffix",
batch_size=64,
progress_bar=False,
normalize_embeddings=True,
meta_fields_to_embed=["meta_field"],
embedding_separator=" - ",
trust_remote_code=True,
local_files_only=True,
truncate_dim=256,
model_kwargs={"torch_dtype": torch.float32},
tokenizer_kwargs={"model_max_length": 512},
config_kwargs={"use_memory_efficient_attention": True},
precision="int8",
encode_kwargs={"task": "clustering"},
)
data = component.to_dict()
assert data == {
"type": "haystack.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", # noqa: E501
"init_parameters": {
"model": "model",
"device": ComponentDevice.from_str("cuda:0").to_dict(),
"token": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
"prefix": "prefix",
"suffix": "suffix",
"batch_size": 64,
"progress_bar": False,
"normalize_embeddings": True,
"embedding_separator": " - ",
"trust_remote_code": True,
"revision": None,
"local_files_only": True,
"meta_fields_to_embed": ["meta_field"],
"truncate_dim": 256,
"model_kwargs": {"torch_dtype": "torch.float32"},
"tokenizer_kwargs": {"model_max_length": 512},
"config_kwargs": {"use_memory_efficient_attention": True},
"precision": "int8",
"encode_kwargs": {"task": "clustering"},
"backend": "torch",
},
}
def test_from_dict(self):
init_parameters = {
"model": "model",
"device": ComponentDevice.from_str("cuda:0").to_dict(),
"token": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
"prefix": "prefix",
"suffix": "suffix",
"batch_size": 64,
"progress_bar": False,
"normalize_embeddings": True,
"embedding_separator": " - ",
"meta_fields_to_embed": ["meta_field"],
"trust_remote_code": True,
"revision": "v1.0",
"local_files_only": True,
"truncate_dim": 256,
"model_kwargs": {"torch_dtype": "torch.float32"},
"tokenizer_kwargs": {"model_max_length": 512},
"config_kwargs": {"use_memory_efficient_attention": True},
"precision": "int8",
}
component = SentenceTransformersDocumentEmbedder.from_dict(
{
"type": "haystack.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", # noqa: E501
"init_parameters": init_parameters,
}
)
assert component.model == "model"
assert component.device == ComponentDevice.from_str("cuda:0")
assert component.token == Secret.from_env_var("ENV_VAR", strict=False)
assert component.prefix == "prefix"
assert component.suffix == "suffix"
assert component.batch_size == 64
assert component.progress_bar is False
assert component.normalize_embeddings is True
assert component.embedding_separator == " - "
assert component.trust_remote_code
assert component.revision == "v1.0"
assert component.local_files_only
assert component.meta_fields_to_embed == ["meta_field"]
assert component.truncate_dim == 256
assert component.model_kwargs == {"torch_dtype": torch.float32}
assert component.tokenizer_kwargs == {"model_max_length": 512}
assert component.config_kwargs == {"use_memory_efficient_attention": True}
assert component.precision == "int8"
def test_from_dict_no_default_parameters(self):
component = SentenceTransformersDocumentEmbedder.from_dict(
{
"type": "haystack.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", # noqa: E501
"init_parameters": {},
}
)
assert component.model == "sentence-transformers/all-mpnet-base-v2"
assert component.device == ComponentDevice.resolve_device(None)
assert component.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert component.prefix == ""
assert component.suffix == ""
assert component.batch_size == 32
assert component.progress_bar is True
assert component.normalize_embeddings is False
assert component.embedding_separator == "\n"
assert component.trust_remote_code is False
assert component.revision is None
assert component.local_files_only is False
assert component.meta_fields_to_embed == []
assert component.truncate_dim is None
assert component.precision == "float32"
def test_from_dict_none_device(self):
init_parameters = {
"model": "model",
"device": None,
"token": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
"prefix": "prefix",
"suffix": "suffix",
"batch_size": 64,
"progress_bar": False,
"normalize_embeddings": True,
"embedding_separator": " - ",
"meta_fields_to_embed": ["meta_field"],
"trust_remote_code": True,
"local_files_only": False,
"truncate_dim": None,
"precision": "float32",
}
component = SentenceTransformersDocumentEmbedder.from_dict(
{
"type": "haystack.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder", # noqa: E501
"init_parameters": init_parameters,
}
)
assert component.model == "model"
assert component.device == ComponentDevice.resolve_device(None)
assert component.token == Secret.from_env_var("ENV_VAR", strict=False)
assert component.prefix == "prefix"
assert component.suffix == "suffix"
assert component.batch_size == 64
assert component.progress_bar is False
assert component.normalize_embeddings is True
assert component.embedding_separator == " - "
assert component.trust_remote_code
assert component.revision is None
assert component.local_files_only is False
assert component.meta_fields_to_embed == ["meta_field"]
assert component.truncate_dim is None
assert component.precision == "float32"
@patch(
"haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory"
)
def test_warmup(self, mocked_factory):
embedder = SentenceTransformersDocumentEmbedder(
model="model",
token=None,
device=ComponentDevice.from_str("cpu"),
tokenizer_kwargs={"model_max_length": 512},
config_kwargs={"use_memory_efficient_attention": True},
)
mocked_factory.get_embedding_backend.assert_not_called()
embedder.warm_up()
embedder.embedding_backend.model.max_seq_length = 512
mocked_factory.get_embedding_backend.assert_called_once_with(
model="model",
device="cpu",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
truncate_dim=None,
model_kwargs=None,
tokenizer_kwargs={"model_max_length": 512},
config_kwargs={"use_memory_efficient_attention": True},
backend="torch",
)
@patch(
"haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory"
)
def test_warmup_doesnt_reload(self, mocked_factory):
embedder = SentenceTransformersDocumentEmbedder(model="model")
mocked_factory.get_embedding_backend.assert_not_called()
embedder.warm_up()
embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once()
def test_run(self):
embedder = SentenceTransformersDocumentEmbedder(model="model")
embedder.embedding_backend = MagicMock()
embedder.embedding_backend.embed = lambda x, **_: [[random.random() for _ in range(16)] for _ in range(len(x))]
documents = [Document(content=f"document number {i}") for i in range(5)]
result = embedder.run(documents=documents)
assert isinstance(result["documents"], list)
assert len(result["documents"]) == len(documents)
for doc, new_doc in zip(documents, result["documents"], strict=True):
assert new_doc is not doc
assert doc.embedding is None
assert isinstance(new_doc, Document)
assert isinstance(new_doc.embedding, list)
assert isinstance(new_doc.embedding[0], float)
def test_run_wrong_input_format(self):
embedder = SentenceTransformersDocumentEmbedder(model="model")
string_input = "text"
list_integers_input = [1, 2, 3]
with pytest.raises(
TypeError, match="SentenceTransformersDocumentEmbedder expects a list of Documents as input"
):
embedder.run(documents=string_input)
with pytest.raises(
TypeError, match="SentenceTransformersDocumentEmbedder expects a list of Documents as input"
):
embedder.run(documents=list_integers_input)
def test_embed_metadata(self):
embedder = SentenceTransformersDocumentEmbedder(
model="model", meta_fields_to_embed=["meta_field"], embedding_separator="\n"
)
embedder.embedding_backend = MagicMock()
embedder.embedding_backend.embed.return_value = [[random.random() for _ in range(16)] for _ in range(5)]
documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)]
embedder.run(documents=documents)
embedder.embedding_backend.embed.assert_called_once_with(
[
"meta_value 0\ndocument number 0",
"meta_value 1\ndocument number 1",
"meta_value 2\ndocument number 2",
"meta_value 3\ndocument number 3",
"meta_value 4\ndocument number 4",
],
batch_size=32,
show_progress_bar=True,
normalize_embeddings=False,
precision="float32",
)
def test_embed_metadata_preserves_falsy_values(self):
embedder = SentenceTransformersDocumentEmbedder(
model="model",
meta_fields_to_embed=["rating", "is_available", "missing", "none_value"],
embedding_separator="\n",
)
embedder.embedding_backend = MagicMock()
embedder.embedding_backend.embed.return_value = [[random.random() for _ in range(16)]]
documents = [Document(content="document", meta={"rating": 0, "is_available": False, "none_value": None})]
embedder.run(documents=documents)
embedder.embedding_backend.embed.assert_called_once_with(
["0\nFalse\ndocument"],
batch_size=32,
show_progress_bar=True,
normalize_embeddings=False,
precision="float32",
)
def test_embed_encode_kwargs(self):
embedder = SentenceTransformersDocumentEmbedder(model="model", encode_kwargs={"task": "retrieval.passage"})
embedder.embedding_backend = MagicMock()
embedder.embedding_backend.embed.return_value = [[random.random() for _ in range(16)] for _ in range(5)]
documents = [Document(content=f"document number {i}") for i in range(5)]
embedder.run(documents=documents)
embedder.embedding_backend.embed.assert_called_once_with(
["document number 0", "document number 1", "document number 2", "document number 3", "document number 4"],
batch_size=32,
show_progress_bar=True,
normalize_embeddings=False,
precision="float32",
task="retrieval.passage",
)
def test_prefix_suffix(self):
embedder = SentenceTransformersDocumentEmbedder(
model="model",
prefix="my_prefix ",
suffix=" my_suffix",
meta_fields_to_embed=["meta_field"],
embedding_separator="\n",
)
embedder.embedding_backend = MagicMock()
embedder.embedding_backend.embed.return_value = [[random.random() for _ in range(16)] for _ in range(5)]
documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)]
embedder.run(documents=documents)
embedder.embedding_backend.embed.assert_called_once_with(
[
"my_prefix meta_value 0\ndocument number 0 my_suffix",
"my_prefix meta_value 1\ndocument number 1 my_suffix",
"my_prefix meta_value 2\ndocument number 2 my_suffix",
"my_prefix meta_value 3\ndocument number 3 my_suffix",
"my_prefix meta_value 4\ndocument number 4 my_suffix",
],
batch_size=32,
show_progress_bar=True,
normalize_embeddings=False,
precision="float32",
)
@patch(
"haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory"
)
def test_model_onnx_backend(self, mocked_factory):
onnx_embedder = SentenceTransformersDocumentEmbedder(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cpu"),
# setting the path isn't necessary if the repo contains a "onnx/model.onnx" file but this is to prevent
# a HF warning
model_kwargs={"file_name": "onnx/model.onnx"},
backend="onnx",
)
onnx_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="sentence-transformers/all-MiniLM-L6-v2",
device="cpu",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
truncate_dim=None,
model_kwargs={"file_name": "onnx/model.onnx"},
tokenizer_kwargs=None,
config_kwargs=None,
backend="onnx",
)
@patch(
"haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory"
)
def test_model_openvino_backend(self, mocked_factory):
openvino_embedder = SentenceTransformersDocumentEmbedder(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cpu"),
# setting the path isn't necessary if the repo contains a "openvino/openvino_model.xml" file but this is
# to prevent a HF warning
model_kwargs={"file_name": "openvino/openvino_model.xml"},
backend="openvino",
)
openvino_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="sentence-transformers/all-MiniLM-L6-v2",
device="cpu",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
truncate_dim=None,
model_kwargs={"file_name": "openvino/openvino_model.xml"},
tokenizer_kwargs=None,
config_kwargs=None,
backend="openvino",
)
@patch(
"haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory"
)
@pytest.mark.parametrize("model_kwargs", [{"torch_dtype": "bfloat16"}, {"torch_dtype": "float16"}])
def test_dtype_on_gpu(self, mocked_factory, model_kwargs):
torch_dtype_embedder = SentenceTransformersDocumentEmbedder(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cuda:0"),
model_kwargs=model_kwargs,
)
torch_dtype_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="sentence-transformers/all-MiniLM-L6-v2",
device="cuda:0",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
truncate_dim=None,
model_kwargs=model_kwargs,
tokenizer_kwargs=None,
config_kwargs=None,
backend="torch",
)
test/components/embedders/test_sentence_transformers_embedding_backend.py
-66
View File
@@ -1,66 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from unittest.mock import patch
from haystack.components.embedders.backends.sentence_transformers_backend import (
_SentenceTransformersEmbeddingBackendFactory,
)
from haystack.utils.auth import Secret
@patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer")
def test_factory_behavior(mock_sentence_transformer):
embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model="my_model", device="cpu"
)
same_embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model="my_model", device="cpu"
)
another_embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model="another_model", device="cpu"
)
yet_another_embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model="my_model", device="cpu", trust_remote_code=True
)
assert same_embedding_backend is embedding_backend
assert another_embedding_backend is not embedding_backend
assert yet_another_embedding_backend is not embedding_backend
@patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer")
def test_model_initialization(mock_sentence_transformer):
_SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model="model",
device="cpu",
auth_token=Secret.from_token("fake-api-token"),
trust_remote_code=True,
local_files_only=True,
truncate_dim=256,
backend="torch",
)
mock_sentence_transformer.assert_called_once_with(
model_name_or_path="model",
device="cpu",
token="fake-api-token",
trust_remote_code=True,
revision=None,
local_files_only=True,
truncate_dim=256,
model_kwargs=None,
tokenizer_kwargs=None,
config_kwargs=None,
backend="torch",
)
@patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer")
def test_embedding_function_with_kwargs(mock_sentence_transformer):
embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(model="model")
data = ["sentence1", "sentence2"]
embedding_backend.embed(data=data, normalize_embeddings=True)
embedding_backend.model.encode.assert_called_once_with(data, normalize_embeddings=True)
test/components/embedders/test_sentence_transformers_sparse_document_embedder.py
-465
View File
@@ -1,465 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from unittest.mock import MagicMock, patch
import pytest
import torch
from haystack import Document
from haystack.components.embedders.sentence_transformers_sparse_document_embedder import (
SentenceTransformersSparseDocumentEmbedder,
)
from haystack.dataclasses.sparse_embedding import SparseEmbedding
from haystack.utils import ComponentDevice, Secret
TYPE_NAME = (
"haystack.components.embedders.sentence_transformers_sparse_document_embedder."
"SentenceTransformersSparseDocumentEmbedder"
)
class TestSentenceTransformersDocumentEmbedder:
def test_init_default(self):
embedder = SentenceTransformersSparseDocumentEmbedder(model="model")
assert embedder.model == "model"
assert embedder.device == ComponentDevice.resolve_device(None)
assert embedder.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert embedder.prefix == ""
assert embedder.suffix == ""
assert embedder.batch_size == 32
assert embedder.progress_bar is True
assert embedder.meta_fields_to_embed == []
assert embedder.embedding_separator == "\n"
assert embedder.trust_remote_code is False
assert embedder.revision is None
assert embedder.local_files_only is False
def test_init_with_parameters(self):
embedder = SentenceTransformersSparseDocumentEmbedder(
model="model",
device=ComponentDevice.from_str("cuda:0"),
token=Secret.from_token("fake-api-token"),
prefix="prefix",
suffix="suffix",
batch_size=64,
progress_bar=False,
meta_fields_to_embed=["test_field"],
embedding_separator=" | ",
trust_remote_code=True,
revision="v1.0",
local_files_only=True,
)
assert embedder.model == "model"
assert embedder.device == ComponentDevice.from_str("cuda:0")
assert embedder.token == Secret.from_token("fake-api-token")
assert embedder.prefix == "prefix"
assert embedder.suffix == "suffix"
assert embedder.batch_size == 64
assert embedder.progress_bar is False
assert embedder.meta_fields_to_embed == ["test_field"]
assert embedder.embedding_separator == " | "
assert embedder.trust_remote_code
assert embedder.revision == "v1.0"
assert embedder.local_files_only
def test_to_dict(self):
component = SentenceTransformersSparseDocumentEmbedder(model="model", device=ComponentDevice.from_str("cpu"))
data = component.to_dict()
assert data == {
"type": TYPE_NAME,
"init_parameters": {
"model": "model",
"device": ComponentDevice.from_str("cpu").to_dict(),
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"prefix": "",
"suffix": "",
"batch_size": 32,
"progress_bar": True,
"embedding_separator": "\n",
"meta_fields_to_embed": [],
"trust_remote_code": False,
"revision": None,
"local_files_only": False,
"model_kwargs": None,
"tokenizer_kwargs": None,
"config_kwargs": None,
"backend": "torch",
},
}
def test_to_dict_with_custom_init_parameters(self):
component = SentenceTransformersSparseDocumentEmbedder(
model="model",
device=ComponentDevice.from_str("cuda:0"),
token=Secret.from_env_var("ENV_VAR", strict=False),
prefix="prefix",
suffix="suffix",
batch_size=64,
progress_bar=False,
meta_fields_to_embed=["meta_field"],
embedding_separator=" - ",
trust_remote_code=True,
local_files_only=True,
model_kwargs={"torch_dtype": torch.float32},
tokenizer_kwargs={"model_max_length": 512},
config_kwargs={"use_memory_efficient_attention": True},
)
data = component.to_dict()
assert data == {
"type": TYPE_NAME,
"init_parameters": {
"model": "model",
"device": ComponentDevice.from_str("cuda:0").to_dict(),
"token": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
"prefix": "prefix",
"suffix": "suffix",
"batch_size": 64,
"progress_bar": False,
"embedding_separator": " - ",
"trust_remote_code": True,
"revision": None,
"local_files_only": True,
"meta_fields_to_embed": ["meta_field"],
"model_kwargs": {"torch_dtype": "torch.float32"},
"tokenizer_kwargs": {"model_max_length": 512},
"config_kwargs": {"use_memory_efficient_attention": True},
"backend": "torch",
},
}
def test_from_dict(self):
init_parameters = {
"model": "model",
"device": ComponentDevice.from_str("cuda:0").to_dict(),
"token": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
"prefix": "prefix",
"suffix": "suffix",
"batch_size": 64,
"progress_bar": False,
"embedding_separator": " - ",
"meta_fields_to_embed": ["meta_field"],
"trust_remote_code": True,
"revision": "v1.0",
"local_files_only": True,
"model_kwargs": {"torch_dtype": "torch.float32"},
"tokenizer_kwargs": {"model_max_length": 512},
"config_kwargs": {"use_memory_efficient_attention": True},
}
component = SentenceTransformersSparseDocumentEmbedder.from_dict(
{"type": TYPE_NAME, "init_parameters": init_parameters}
)
assert component.model == "model"
assert component.device == ComponentDevice.from_str("cuda:0")
assert component.token == Secret.from_env_var("ENV_VAR", strict=False)
assert component.prefix == "prefix"
assert component.suffix == "suffix"
assert component.batch_size == 64
assert component.progress_bar is False
assert component.embedding_separator == " - "
assert component.trust_remote_code
assert component.revision == "v1.0"
assert component.local_files_only
assert component.meta_fields_to_embed == ["meta_field"]
assert component.model_kwargs == {"torch_dtype": torch.float32}
assert component.tokenizer_kwargs == {"model_max_length": 512}
assert component.config_kwargs == {"use_memory_efficient_attention": True}
def test_from_dict_no_default_parameters(self):
component = SentenceTransformersSparseDocumentEmbedder.from_dict({"type": TYPE_NAME, "init_parameters": {}})
assert component.model == "prithivida/Splade_PP_en_v2"
assert component.device == ComponentDevice.resolve_device(None)
assert component.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert component.prefix == ""
assert component.suffix == ""
assert component.batch_size == 32
assert component.progress_bar is True
assert component.embedding_separator == "\n"
assert component.trust_remote_code is False
assert component.revision is None
assert component.local_files_only is False
assert component.meta_fields_to_embed == []
def test_from_dict_none_device(self):
init_parameters = {
"model": "model",
"device": None,
"token": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
"prefix": "prefix",
"suffix": "suffix",
"batch_size": 64,
"progress_bar": False,
"embedding_separator": " - ",
"meta_fields_to_embed": ["meta_field"],
"trust_remote_code": True,
"local_files_only": False,
}
component = SentenceTransformersSparseDocumentEmbedder.from_dict(
{"type": TYPE_NAME, "init_parameters": init_parameters}
)
assert component.model == "model"
assert component.device == ComponentDevice.resolve_device(None)
assert component.token == Secret.from_env_var("ENV_VAR", strict=False)
assert component.prefix == "prefix"
assert component.suffix == "suffix"
assert component.batch_size == 64
assert component.progress_bar is False
assert component.embedding_separator == " - "
assert component.trust_remote_code
assert component.revision is None
assert component.local_files_only is False
assert component.meta_fields_to_embed == ["meta_field"]
@patch(
"haystack.components.embedders.sentence_transformers_sparse_document_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
)
def test_warmup(self, mocked_factory):
embedder = SentenceTransformersSparseDocumentEmbedder(
model="model",
token=None,
device=ComponentDevice.from_str("cpu"),
tokenizer_kwargs={"model_max_length": 512},
config_kwargs={"use_memory_efficient_attention": True},
)
mocked_factory.get_embedding_backend.assert_not_called()
embedder.warm_up()
embedder.embedding_backend.model.max_seq_length = 512
mocked_factory.get_embedding_backend.assert_called_once_with(
model="model",
device="cpu",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
model_kwargs=None,
tokenizer_kwargs={"model_max_length": 512},
config_kwargs={"use_memory_efficient_attention": True},
backend="torch",
)
@patch(
"haystack.components.embedders.sentence_transformers_sparse_document_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
)
def test_warmup_doesnt_reload(self, mocked_factory):
embedder = SentenceTransformersSparseDocumentEmbedder(model="model")
mocked_factory.get_embedding_backend.assert_not_called()
embedder.warm_up()
embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once()
def test_run(self):
embedder = SentenceTransformersSparseDocumentEmbedder(model="model")
embedder.embedding_backend = MagicMock()
def fake_embed(data, **kwargs):
return [SparseEmbedding(indices=[0, 2, 5], values=[0.1, 0.2, 0.3]) for _ in range(len(data))]
embedder.embedding_backend.embed = fake_embed
documents = [Document(content=f"document number {i}") for i in range(5)]
result = embedder.run(documents=documents)
assert isinstance(result["documents"], list)
assert len(result["documents"]) == len(documents)
for doc in result["documents"]:
assert isinstance(doc, Document)
assert isinstance(doc.sparse_embedding, SparseEmbedding)
assert isinstance(doc.sparse_embedding.indices[0], int)
assert isinstance(doc.sparse_embedding.values[0], float)
def test_run_wrong_input_format(self):
embedder = SentenceTransformersSparseDocumentEmbedder(model="model")
string_input = "text"
list_integers_input = [1, 2, 3]
with pytest.raises(
TypeError, match="SentenceTransformersSparseDocumentEmbedder expects a list of Documents as input"
):
embedder.run(documents=string_input)
with pytest.raises(
TypeError, match="SentenceTransformersSparseDocumentEmbedder expects a list of Documents as input"
):
embedder.run(documents=list_integers_input)
def test_embed_metadata(self):
embedder = SentenceTransformersSparseDocumentEmbedder(
model="model", meta_fields_to_embed=["meta_field"], embedding_separator="\n"
)
embedder.embedding_backend = MagicMock()
embedder.embedding_backend.embed.return_value = [
SparseEmbedding(indices=[0, 2, 5], values=[0.1, 0.2, 0.3]) for _ in range(5)
]
documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)]
embedder.run(documents=documents)
embedder.embedding_backend.embed.assert_called_once_with(
data=[
"meta_value 0\ndocument number 0",
"meta_value 1\ndocument number 1",
"meta_value 2\ndocument number 2",
"meta_value 3\ndocument number 3",
"meta_value 4\ndocument number 4",
],
batch_size=32,
show_progress_bar=True,
)
def test_embed_metadata_preserves_falsy_values(self):
embedder = SentenceTransformersSparseDocumentEmbedder(
model="model",
meta_fields_to_embed=["rating", "is_available", "missing", "none_value"],
embedding_separator="\n",
)
embedder.embedding_backend = MagicMock()
embedder.embedding_backend.embed.return_value = [SparseEmbedding(indices=[0], values=[0.1])]
documents = [Document(content="document", meta={"rating": 0, "is_available": False, "none_value": None})]
embedder.run(documents=documents)
embedder.embedding_backend.embed.assert_called_once_with(
data=["0\nFalse\ndocument"], batch_size=32, show_progress_bar=True
)
def test_prefix_suffix(self):
embedder = SentenceTransformersSparseDocumentEmbedder(
model="model",
prefix="my_prefix ",
suffix=" my_suffix",
meta_fields_to_embed=["meta_field"],
embedding_separator="\n",
)
embedder.embedding_backend = MagicMock()
embedder.embedding_backend.embed.return_value = [
SparseEmbedding(indices=[0, 2, 5], values=[0.1, 0.2, 0.3]) for _ in range(5)
]
documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)]
embedder.run(documents=documents)
embedder.embedding_backend.embed.assert_called_once_with(
data=[
"my_prefix meta_value 0\ndocument number 0 my_suffix",
"my_prefix meta_value 1\ndocument number 1 my_suffix",
"my_prefix meta_value 2\ndocument number 2 my_suffix",
"my_prefix meta_value 3\ndocument number 3 my_suffix",
"my_prefix meta_value 4\ndocument number 4 my_suffix",
],
batch_size=32,
show_progress_bar=True,
)
@patch(
"haystack.components.embedders.sentence_transformers_sparse_document_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
)
def test_model_onnx_backend(self, mocked_factory):
onnx_embedder = SentenceTransformersSparseDocumentEmbedder(
model="prithivida/Splade_PP_en_v2",
token=None,
device=ComponentDevice.from_str("cpu"),
model_kwargs={
"file_name": "onnx/model.onnx"
}, # setting the path isn't necessary if the repo contains a "onnx/model.onnx" file but this is to
# prevent a HF warning
backend="onnx",
)
onnx_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="prithivida/Splade_PP_en_v2",
device="cpu",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
model_kwargs={"file_name": "onnx/model.onnx"},
tokenizer_kwargs=None,
config_kwargs=None,
backend="onnx",
)
@patch(
"haystack.components.embedders.sentence_transformers_sparse_document_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
)
def test_model_openvino_backend(self, mocked_factory):
openvino_embedder = SentenceTransformersSparseDocumentEmbedder(
model="prithivida/Splade_PP_en_v2",
token=None,
device=ComponentDevice.from_str("cpu"),
model_kwargs={
"file_name": "openvino/openvino_model.xml"
}, # setting the path isn't necessary if the repo contains a "openvino/openvino_model.xml" file but this
# is to prevent a HF warning
backend="openvino",
)
openvino_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="prithivida/Splade_PP_en_v2",
device="cpu",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
model_kwargs={"file_name": "openvino/openvino_model.xml"},
tokenizer_kwargs=None,
config_kwargs=None,
backend="openvino",
)
@patch(
"haystack.components.embedders.sentence_transformers_sparse_document_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
)
@pytest.mark.parametrize("model_kwargs", [{"torch_dtype": "bfloat16"}, {"torch_dtype": "float16"}])
def test_dtype_on_gpu(self, mocked_factory, model_kwargs):
torch_dtype_embedder = SentenceTransformersSparseDocumentEmbedder(
model="prithivida/Splade_PP_en_v2",
token=None,
device=ComponentDevice.from_str("cuda:0"),
model_kwargs=model_kwargs,
)
torch_dtype_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="prithivida/Splade_PP_en_v2",
device="cuda:0",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
model_kwargs=model_kwargs,
tokenizer_kwargs=None,
config_kwargs=None,
backend="torch",
)
@pytest.mark.integration
@pytest.mark.slow
@pytest.mark.flaky(reruns=3, reruns_delay=10)
def test_live_run_sparse_document_embedder(self, del_hf_env_vars):
docs = [
Document(content="I love cheese", meta={"topic": "Cuisine"}),
Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}),
]
embedder = SentenceTransformersSparseDocumentEmbedder(
model="sparse-encoder-testing/splade-bert-tiny-nq",
meta_fields_to_embed=["topic"],
embedding_separator=" | ",
device=ComponentDevice.from_str("cpu"),
)
result = embedder.run(documents=docs)
documents_with_embeddings = result["documents"]
assert isinstance(documents_with_embeddings, list)
assert len(documents_with_embeddings) == len(docs)
for doc in documents_with_embeddings:
assert isinstance(doc, Document)
assert hasattr(doc, "sparse_embedding")
assert isinstance(doc.sparse_embedding, SparseEmbedding)
assert isinstance(doc.sparse_embedding.indices, list)
assert isinstance(doc.sparse_embedding.values, list)
assert len(doc.sparse_embedding.indices) == len(doc.sparse_embedding.values)
# Expect at least one non-zero entry
assert len(doc.sparse_embedding.indices) > 0
test/components/embedders/test_sentence_transformers_sparse_embedding_backend.py
-104
View File
@@ -1,104 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from unittest.mock import patch
import pytest
import torch
from haystack.components.embedders.backends.sentence_transformers_sparse_backend import (
_SentenceTransformersSparseEmbeddingBackendFactory,
)
from haystack.dataclasses.sparse_embedding import SparseEmbedding
from haystack.utils.auth import Secret
@patch("haystack.components.embedders.backends.sentence_transformers_sparse_backend.SparseEncoder")
def test_sparse_factory_behavior(mock_sparse_encoder):
embedding_backend = _SentenceTransformersSparseEmbeddingBackendFactory.get_embedding_backend(
model="my_model", device="cpu"
)
same_embedding_backend = _SentenceTransformersSparseEmbeddingBackendFactory.get_embedding_backend(
model="my_model", device="cpu"
)
another_embedding_backend = _SentenceTransformersSparseEmbeddingBackendFactory.get_embedding_backend(
model="another_model", device="cpu"
)
yet_another_embedding_backend = _SentenceTransformersSparseEmbeddingBackendFactory.get_embedding_backend(
model="my_model", device="cpu", trust_remote_code=True
)
assert same_embedding_backend is embedding_backend
assert another_embedding_backend is not embedding_backend
assert yet_another_embedding_backend is not embedding_backend
@patch("haystack.components.embedders.backends.sentence_transformers_sparse_backend.SparseEncoder")
def test_sparse_model_initialization(mock_sparse_encoder):
_SentenceTransformersSparseEmbeddingBackendFactory.get_embedding_backend(
model="model",
device="cpu",
auth_token=Secret.from_token("fake-api-token"),
trust_remote_code=True,
local_files_only=True,
backend="torch",
)
mock_sparse_encoder.assert_called_once_with(
model_name_or_path="model",
device="cpu",
token="fake-api-token",
trust_remote_code=True,
revision=None,
local_files_only=True,
model_kwargs=None,
tokenizer_kwargs=None,
config_kwargs=None,
backend="torch",
)
@patch("haystack.components.embedders.backends.sentence_transformers_sparse_backend.SparseEncoder")
def test_sparse_embedding_function_with_kwargs(mock_sparse_encoder):
embedding_backend = _SentenceTransformersSparseEmbeddingBackendFactory.get_embedding_backend(model="model")
data = ["sentence1", "sentence2"]
embedding_backend.embed(data=data, attn_implementation="sdpa")
embedding_backend.model.encode.assert_called_once_with(
data, convert_to_tensor=False, convert_to_sparse_tensor=True, attn_implementation="sdpa"
)
@patch("haystack.components.embedders.backends.sentence_transformers_sparse_backend.SparseEncoder")
def test_sparse_embedding_function(mock_sparse_encoder):
"""
Test that the backend's embed method returns the correct sparse embeddings.
"""
# Ensure the factory cache is cleared before each test.
_SentenceTransformersSparseEmbeddingBackendFactory._instances = {}
tensors = [
torch.sparse_coo_tensor(torch.tensor([[1, 4]]), torch.tensor([0.5, 0.8])),
torch.sparse_coo_tensor(torch.tensor([[2]]), torch.tensor([0.3])),
]
mock_sparse_encoder.return_value.encode.return_value = tensors
# Get the embedding backend
embedding_backend = _SentenceTransformersSparseEmbeddingBackendFactory.get_embedding_backend(model="model")
# Embed dummy data
data = ["sentence1", "sentence2"]
sparse_embeddings = embedding_backend.embed(data=data)
# Expected output
expected_embeddings = [
SparseEmbedding(indices=[1, 4], values=[0.5, 0.8]),
SparseEmbedding(indices=[2], values=[0.3]),
]
assert len(sparse_embeddings) == len(expected_embeddings)
for got, exp in zip(sparse_embeddings, expected_embeddings, strict=True):
assert got.indices == exp.indices
assert got.values == pytest.approx(exp.values)
test/components/embedders/test_sentence_transformers_sparse_text_embedder.py
-339
View File
@@ -1,339 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from unittest.mock import MagicMock, patch
import pytest
import torch
from haystack.components.embedders.sentence_transformers_sparse_text_embedder import (
SentenceTransformersSparseTextEmbedder,
)
from haystack.dataclasses.sparse_embedding import SparseEmbedding
from haystack.utils import ComponentDevice, Secret
TYPE_NAME = (
"haystack.components.embedders.sentence_transformers_sparse_text_embedder.SentenceTransformersSparseTextEmbedder"
)
class TestSentenceTransformersSparseTextEmbedder:
def test_init_default(self):
embedder = SentenceTransformersSparseTextEmbedder(model="model")
assert embedder.model == "model"
assert embedder.device == ComponentDevice.resolve_device(None)
assert embedder.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert embedder.prefix == ""
assert embedder.suffix == ""
assert embedder.trust_remote_code is False
assert embedder.revision is None
assert embedder.local_files_only is False
def test_init_with_parameters(self):
embedder = SentenceTransformersSparseTextEmbedder(
model="model",
device=ComponentDevice.from_str("cuda:0"),
token=Secret.from_token("fake-api-token"),
prefix="prefix",
suffix="suffix",
trust_remote_code=True,
revision="v1.0",
local_files_only=True,
)
assert embedder.model == "model"
assert embedder.device == ComponentDevice.from_str("cuda:0")
assert embedder.token == Secret.from_token("fake-api-token")
assert embedder.prefix == "prefix"
assert embedder.suffix == "suffix"
assert embedder.trust_remote_code is True
assert embedder.revision == "v1.0"
assert embedder.local_files_only is True
def test_to_dict(self):
component = SentenceTransformersSparseTextEmbedder(model="model", device=ComponentDevice.from_str("cpu"))
data = component.to_dict()
assert data == {
"type": TYPE_NAME,
"init_parameters": {
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"model": "model",
"device": ComponentDevice.from_str("cpu").to_dict(),
"prefix": "",
"suffix": "",
"trust_remote_code": False,
"revision": None,
"local_files_only": False,
"model_kwargs": None,
"tokenizer_kwargs": None,
"config_kwargs": None,
"backend": "torch",
},
}
def test_to_dict_with_custom_init_parameters(self):
component = SentenceTransformersSparseTextEmbedder(
model="model",
device=ComponentDevice.from_str("cuda:0"),
token=Secret.from_env_var("ENV_VAR", strict=False),
prefix="prefix",
suffix="suffix",
trust_remote_code=True,
local_files_only=True,
model_kwargs={"torch_dtype": torch.float32},
tokenizer_kwargs={"model_max_length": 512},
config_kwargs={"use_memory_efficient_attention": False},
)
data = component.to_dict()
assert data == {
"type": TYPE_NAME,
"init_parameters": {
"token": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
"model": "model",
"device": ComponentDevice.from_str("cuda:0").to_dict(),
"prefix": "prefix",
"suffix": "suffix",
"trust_remote_code": True,
"revision": None,
"local_files_only": True,
"model_kwargs": {"torch_dtype": "torch.float32"},
"tokenizer_kwargs": {"model_max_length": 512},
"config_kwargs": {"use_memory_efficient_attention": False},
"backend": "torch",
},
}
def test_to_dict_not_serialize_token(self):
component = SentenceTransformersSparseTextEmbedder(model="model", token=Secret.from_token("fake-api-token"))
with pytest.raises(ValueError, match="Cannot serialize token-based secret"):
component.to_dict()
def test_from_dict(self):
data = {
"type": TYPE_NAME,
"init_parameters": {
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"model": "model",
"device": ComponentDevice.from_str("cpu").to_dict(),
"prefix": "",
"suffix": "",
"trust_remote_code": False,
"revision": "v1.0",
"local_files_only": False,
"model_kwargs": {"torch_dtype": "torch.float32"},
"tokenizer_kwargs": {"model_max_length": 512},
"config_kwargs": {"use_memory_efficient_attention": False},
},
}
component = SentenceTransformersSparseTextEmbedder.from_dict(data)
assert component.model == "model"
assert component.device == ComponentDevice.from_str("cpu")
assert component.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert component.prefix == ""
assert component.suffix == ""
assert component.trust_remote_code is False
assert component.revision == "v1.0"
assert component.local_files_only is False
assert component.model_kwargs == {"torch_dtype": torch.float32}
assert component.tokenizer_kwargs == {"model_max_length": 512}
assert component.config_kwargs == {"use_memory_efficient_attention": False}
def test_from_dict_no_default_parameters(self):
data = {"type": TYPE_NAME, "init_parameters": {}}
component = SentenceTransformersSparseTextEmbedder.from_dict(data)
assert component.model == "prithivida/Splade_PP_en_v2"
assert component.device == ComponentDevice.resolve_device(None)
assert component.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert component.prefix == ""
assert component.suffix == ""
assert component.trust_remote_code is False
assert component.revision is None
assert component.local_files_only is False
def test_from_dict_none_device(self):
data = {
"type": TYPE_NAME,
"init_parameters": {
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"model": "model",
"device": None,
"prefix": "",
"suffix": "",
"trust_remote_code": False,
"local_files_only": False,
},
}
component = SentenceTransformersSparseTextEmbedder.from_dict(data)
assert component.model == "model"
assert component.device == ComponentDevice.resolve_device(None)
assert component.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert component.prefix == ""
assert component.suffix == ""
assert component.trust_remote_code is False
assert component.revision is None
assert component.local_files_only is False
@patch(
"haystack.components.embedders.sentence_transformers_sparse_text_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
)
def test_warmup(self, mocked_factory):
embedder = SentenceTransformersSparseTextEmbedder(
model="model",
token=None,
device=ComponentDevice.from_str("cpu"),
tokenizer_kwargs={"model_max_length": 512},
)
mocked_factory.get_embedding_backend.assert_not_called()
embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="model",
device="cpu",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
model_kwargs=None,
tokenizer_kwargs={"model_max_length": 512},
config_kwargs=None,
backend="torch",
)
@patch(
"haystack.components.embedders.sentence_transformers_sparse_text_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
)
def test_warmup_doesnt_reload(self, mocked_factory):
embedder = SentenceTransformersSparseTextEmbedder(model="model")
mocked_factory.get_embedding_backend.assert_not_called()
embedder.warm_up()
embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once()
def test_run(self):
embedder = SentenceTransformersSparseTextEmbedder(model="model")
embedder.embedding_backend = MagicMock()
def fake_embed(data, **kwargs):
return [SparseEmbedding(indices=[1, 3], values=[0.5, 0.7]) for _ in range(len(data))]
embedder.embedding_backend.embed = fake_embed
text = "a nice text to embed"
result = embedder.run(text=text)
sparse_embedding = result["sparse_embedding"]
assert isinstance(sparse_embedding, SparseEmbedding)
assert sparse_embedding.indices == [1, 3]
assert sparse_embedding.values == [0.5, 0.7]
def test_run_wrong_input_format(self):
embedder = SentenceTransformersSparseTextEmbedder(model="model")
embedder.embedding_backend = MagicMock()
list_integers_input = [1, 2, 3]
with pytest.raises(TypeError, match="SentenceTransformersSparseTextEmbedder expects a string as input"):
embedder.run(text=list_integers_input)
@patch(
"haystack.components.embedders.sentence_transformers_sparse_text_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
)
def test_model_onnx_backend(self, mocked_factory):
onnx_embedder = SentenceTransformersSparseTextEmbedder(
model="prithivida/Splade_PP_en_v2",
token=None,
device=ComponentDevice.from_str("cpu"),
model_kwargs={
"file_name": "onnx/model.onnx"
}, # setting the path isn't necessary if the repo contains a "onnx/model.onnx" file but this is to
# prevent a HF warning
backend="onnx",
)
onnx_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="prithivida/Splade_PP_en_v2",
device="cpu",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
model_kwargs={"file_name": "onnx/model.onnx"},
tokenizer_kwargs=None,
config_kwargs=None,
backend="onnx",
)
@patch(
"haystack.components.embedders.sentence_transformers_sparse_text_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
)
def test_model_openvino_backend(self, mocked_factory):
openvino_embedder = SentenceTransformersSparseTextEmbedder(
model="prithivida/Splade_PP_en_v2",
token=None,
device=ComponentDevice.from_str("cpu"),
model_kwargs={
"file_name": "openvino/openvino_model.xml"
}, # setting the path isn't necessary if the repo contains a "openvino/openvino_model.xml" file but
# this is to prevent a HF warning
backend="openvino",
)
openvino_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="prithivida/Splade_PP_en_v2",
device="cpu",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
model_kwargs={"file_name": "openvino/openvino_model.xml"},
tokenizer_kwargs=None,
config_kwargs=None,
backend="openvino",
)
@patch(
"haystack.components.embedders.sentence_transformers_sparse_text_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
)
@pytest.mark.parametrize("model_kwargs", [{"torch_dtype": "bfloat16"}, {"torch_dtype": "float16"}])
def test_dtype_on_gpu(self, mocked_factory, model_kwargs):
torch_dtype_embedder = SentenceTransformersSparseTextEmbedder(
model="prithivida/Splade_PP_en_v2",
token=None,
device=ComponentDevice.from_str("cuda:0"),
model_kwargs=model_kwargs,
)
torch_dtype_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="prithivida/Splade_PP_en_v2",
device="cuda:0",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
model_kwargs=model_kwargs,
tokenizer_kwargs=None,
config_kwargs=None,
backend="torch",
)
@pytest.mark.integration
@pytest.mark.slow
@pytest.mark.flaky(reruns=3, reruns_delay=10)
def test_live_run_sparse_text_embedder(self, del_hf_env_vars):
text = "I love Nine Inch Nails"
embedder = SentenceTransformersSparseTextEmbedder(
model="sparse-encoder-testing/splade-bert-tiny-nq", device=ComponentDevice.from_str("cpu")
)
result = embedder.run(text=text)
sparse_embedding = result["sparse_embedding"]
assert isinstance(sparse_embedding, SparseEmbedding)
assert isinstance(sparse_embedding.indices, list)
assert isinstance(sparse_embedding.values, list)
assert len(sparse_embedding.indices) == len(sparse_embedding.values)
assert len(sparse_embedding.indices) > 0
test/components/embedders/test_sentence_transformers_text_embedder.py
-421
View File
@@ -1,421 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
import random
from unittest.mock import MagicMock, patch
import pytest
import torch
from haystack.components.embedders.sentence_transformers_text_embedder import SentenceTransformersTextEmbedder
from haystack.utils import ComponentDevice, Secret
class TestSentenceTransformersTextEmbedder:
def test_init_default(self):
embedder = SentenceTransformersTextEmbedder(model="model")
assert embedder.model == "model"
assert embedder.device == ComponentDevice.resolve_device(None)
assert embedder.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert embedder.prefix == ""
assert embedder.suffix == ""
assert embedder.batch_size == 32
assert embedder.progress_bar is True
assert embedder.normalize_embeddings is False
assert embedder.trust_remote_code is False
assert embedder.revision is None
assert embedder.local_files_only is False
assert embedder.truncate_dim is None
assert embedder.precision == "float32"
def test_init_with_parameters(self):
embedder = SentenceTransformersTextEmbedder(
model="model",
device=ComponentDevice.from_str("cuda:0"),
token=Secret.from_token("fake-api-token"),
prefix="prefix",
suffix="suffix",
batch_size=64,
progress_bar=False,
normalize_embeddings=True,
trust_remote_code=True,
revision="v1.0",
local_files_only=True,
truncate_dim=256,
precision="int8",
)
assert embedder.model == "model"
assert embedder.device == ComponentDevice.from_str("cuda:0")
assert embedder.token == Secret.from_token("fake-api-token")
assert embedder.prefix == "prefix"
assert embedder.suffix == "suffix"
assert embedder.batch_size == 64
assert embedder.progress_bar is False
assert embedder.normalize_embeddings is True
assert embedder.trust_remote_code is True
assert embedder.revision == "v1.0"
assert embedder.local_files_only is True
assert embedder.truncate_dim == 256
assert embedder.precision == "int8"
def test_to_dict(self):
component = SentenceTransformersTextEmbedder(model="model", device=ComponentDevice.from_str("cpu"))
data = component.to_dict()
assert data == {
"type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", # noqa: E501
"init_parameters": {
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"model": "model",
"device": ComponentDevice.from_str("cpu").to_dict(),
"prefix": "",
"suffix": "",
"batch_size": 32,
"progress_bar": True,
"normalize_embeddings": False,
"trust_remote_code": False,
"revision": None,
"local_files_only": False,
"truncate_dim": None,
"model_kwargs": None,
"tokenizer_kwargs": None,
"encode_kwargs": None,
"config_kwargs": None,
"precision": "float32",
"backend": "torch",
},
}
def test_to_dict_with_custom_init_parameters(self):
component = SentenceTransformersTextEmbedder(
model="model",
device=ComponentDevice.from_str("cuda:0"),
token=Secret.from_env_var("ENV_VAR", strict=False),
prefix="prefix",
suffix="suffix",
batch_size=64,
progress_bar=False,
normalize_embeddings=True,
trust_remote_code=True,
local_files_only=True,
truncate_dim=256,
model_kwargs={"torch_dtype": torch.float32},
tokenizer_kwargs={"model_max_length": 512},
config_kwargs={"use_memory_efficient_attention": False},
precision="int8",
encode_kwargs={"task": "clustering"},
)
data = component.to_dict()
assert data == {
"type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", # noqa: E501
"init_parameters": {
"token": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
"model": "model",
"device": ComponentDevice.from_str("cuda:0").to_dict(),
"prefix": "prefix",
"suffix": "suffix",
"batch_size": 64,
"progress_bar": False,
"normalize_embeddings": True,
"trust_remote_code": True,
"revision": None,
"local_files_only": True,
"truncate_dim": 256,
"model_kwargs": {"torch_dtype": "torch.float32"},
"tokenizer_kwargs": {"model_max_length": 512},
"config_kwargs": {"use_memory_efficient_attention": False},
"precision": "int8",
"encode_kwargs": {"task": "clustering"},
"backend": "torch",
},
}
def test_to_dict_not_serialize_token(self):
component = SentenceTransformersTextEmbedder(model="model", token=Secret.from_token("fake-api-token"))
with pytest.raises(ValueError, match="Cannot serialize token-based secret"):
component.to_dict()
def test_from_dict(self):
data = {
"type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", # noqa: E501
"init_parameters": {
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"model": "model",
"device": ComponentDevice.from_str("cpu").to_dict(),
"prefix": "",
"suffix": "",
"batch_size": 32,
"progress_bar": True,
"normalize_embeddings": False,
"trust_remote_code": False,
"revision": "v1.0",
"local_files_only": False,
"truncate_dim": None,
"model_kwargs": {"torch_dtype": "torch.float32"},
"tokenizer_kwargs": {"model_max_length": 512},
"config_kwargs": {"use_memory_efficient_attention": False},
"precision": "float32",
},
}
component = SentenceTransformersTextEmbedder.from_dict(data)
assert component.model == "model"
assert component.device == ComponentDevice.from_str("cpu")
assert component.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert component.prefix == ""
assert component.suffix == ""
assert component.batch_size == 32
assert component.progress_bar is True
assert component.normalize_embeddings is False
assert component.trust_remote_code is False
assert component.revision == "v1.0"
assert component.local_files_only is False
assert component.truncate_dim is None
assert component.model_kwargs == {"torch_dtype": torch.float32}
assert component.tokenizer_kwargs == {"model_max_length": 512}
assert component.config_kwargs == {"use_memory_efficient_attention": False}
assert component.precision == "float32"
def test_from_dict_no_default_parameters(self):
data = {
"type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", # noqa: E501
"init_parameters": {},
}
component = SentenceTransformersTextEmbedder.from_dict(data)
assert component.model == "sentence-transformers/all-mpnet-base-v2"
assert component.device == ComponentDevice.resolve_device(None)
assert component.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert component.prefix == ""
assert component.suffix == ""
assert component.batch_size == 32
assert component.progress_bar is True
assert component.normalize_embeddings is False
assert component.trust_remote_code is False
assert component.revision is None
assert component.local_files_only is False
assert component.truncate_dim is None
assert component.precision == "float32"
def test_from_dict_none_device(self):
data = {
"type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", # noqa: E501
"init_parameters": {
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"model": "model",
"device": None,
"prefix": "",
"suffix": "",
"batch_size": 32,
"progress_bar": True,
"normalize_embeddings": False,
"trust_remote_code": False,
"local_files_only": False,
"truncate_dim": 256,
"precision": "int8",
},
}
component = SentenceTransformersTextEmbedder.from_dict(data)
assert component.model == "model"
assert component.device == ComponentDevice.resolve_device(None)
assert component.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert component.prefix == ""
assert component.suffix == ""
assert component.batch_size == 32
assert component.progress_bar is True
assert component.normalize_embeddings is False
assert component.trust_remote_code is False
assert component.revision is None
assert component.local_files_only is False
assert component.truncate_dim == 256
assert component.precision == "int8"
@patch(
"haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
)
def test_warmup(self, mocked_factory):
embedder = SentenceTransformersTextEmbedder(
model="model",
token=None,
device=ComponentDevice.from_str("cpu"),
tokenizer_kwargs={"model_max_length": 512},
)
mocked_factory.get_embedding_backend.assert_not_called()
embedder.warm_up()
embedder.embedding_backend.model.max_seq_length = 512
mocked_factory.get_embedding_backend.assert_called_once_with(
model="model",
device="cpu",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
truncate_dim=None,
model_kwargs=None,
tokenizer_kwargs={"model_max_length": 512},
config_kwargs=None,
backend="torch",
)
@patch(
"haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
)
def test_warmup_doesnt_reload(self, mocked_factory):
embedder = SentenceTransformersTextEmbedder(model="model")
mocked_factory.get_embedding_backend.assert_not_called()
embedder.warm_up()
embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once()
def test_run(self):
embedder = SentenceTransformersTextEmbedder(model="model")
embedder.embedding_backend = MagicMock()
embedder.embedding_backend.embed = lambda x, **_: [[random.random() for _ in range(16)] for _ in range(len(x))]
text = "a nice text to embed"
result = embedder.run(text=text)
embedding = result["embedding"]
assert isinstance(embedding, list)
assert all(isinstance(el, float) for el in embedding)
def test_run_wrong_input_format(self):
embedder = SentenceTransformersTextEmbedder(model="model")
embedder.embedding_backend = MagicMock()
list_integers_input = [1, 2, 3]
with pytest.raises(TypeError, match="SentenceTransformersTextEmbedder expects a string as input"):
embedder.run(text=list_integers_input)
def test_embed_encode_kwargs(self):
embedder = SentenceTransformersTextEmbedder(model="model", encode_kwargs={"task": "retrieval.query"})
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="float32",
task="retrieval.query",
)
@patch(
"haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
)
def test_model_onnx_backend(self, mocked_factory):
onnx_embedder = SentenceTransformersTextEmbedder(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cpu"),
# setting the path isn't necessary if the repo contains a "onnx/model.onnx" file but this is to prevent
# a HF warning
model_kwargs={"file_name": "onnx/model.onnx"},
backend="onnx",
)
onnx_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="sentence-transformers/all-MiniLM-L6-v2",
device="cpu",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
truncate_dim=None,
model_kwargs={"file_name": "onnx/model.onnx"},
tokenizer_kwargs=None,
config_kwargs=None,
backend="onnx",
)
@patch(
"haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
)
def test_model_openvino_backend(self, mocked_factory):
openvino_embedder = SentenceTransformersTextEmbedder(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cpu"),
# setting the path isn't necessary if the repo contains a "openvino/openvino_model.xml" file but
# this is to prevent a HF warning
model_kwargs={"file_name": "openvino/openvino_model.xml"},
backend="openvino",
)
openvino_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="sentence-transformers/all-MiniLM-L6-v2",
device="cpu",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
truncate_dim=None,
model_kwargs={"file_name": "openvino/openvino_model.xml"},
tokenizer_kwargs=None,
config_kwargs=None,
backend="openvino",
)
@patch(
"haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
)
@pytest.mark.parametrize("model_kwargs", [{"torch_dtype": "bfloat16"}, {"torch_dtype": "float16"}])
def test_dtype_on_gpu(self, mocked_factory, model_kwargs):
torch_dtype_embedder = SentenceTransformersTextEmbedder(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cuda:0"),
model_kwargs=model_kwargs,
)
torch_dtype_embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model="sentence-transformers/all-MiniLM-L6-v2",
device="cuda:0",
auth_token=None,
trust_remote_code=False,
revision=None,
local_files_only=False,
truncate_dim=None,
model_kwargs=model_kwargs,
tokenizer_kwargs=None,
config_kwargs=None,
backend="torch",
)
@pytest.mark.integration
@pytest.mark.slow
def test_run_trunc(self, del_hf_env_vars):
"""
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"
embedder_trunc = SentenceTransformersTextEmbedder(model=checkpoint, truncate_dim=64)
result_trunc = embedder_trunc.run(text=text)
embedding_trunc = result_trunc["embedding"]
# without truncation, the embedding has 128 length
assert len(embedding_trunc) == 64
@pytest.mark.integration
@pytest.mark.slow
def test_run_quantization(self, del_hf_env_vars):
"""
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"
embedder_def = SentenceTransformersTextEmbedder(model=checkpoint, precision="int8")
result_def = embedder_def.run(text=text)
embedding_def = result_def["embedding"]
assert len(embedding_def) == 128
assert all(isinstance(el, int) for el in embedding_def)
test/components/joiners/test_list_joiner.py
+2 -2
View File
@@ -8,7 +8,7 @@ import pytest
from haystack import Document, Pipeline
from haystack.components.builders import AnswerBuilder, ChatPromptBuilder
from haystack.components.embedders import SentenceTransformersTextEmbedder
from haystack.components.embedders import OpenAITextEmbedder
from haystack.components.generators.chat.openai import OpenAIChatGenerator
from haystack.components.joiners.list_joiner import ListJoiner
from haystack.core.errors import PipelineConnectError
@@ -140,7 +140,7 @@ class TestListJoiner:
def test_pipeline_bad_connection(self):
with pytest.raises(PipelineConnectError):
joiner = ListJoiner()
query_embedder = SentenceTransformersTextEmbedder()
query_embedder = OpenAITextEmbedder(api_key=Secret.from_token("test-api-key"))
pipe = Pipeline()
pipe.add_component("joiner", joiner)
pipe.add_component("query_embedder", query_embedder)
test/components/preprocessors/test_embedding_based_document_splitter.py
+16 -25
View File
@@ -9,9 +9,8 @@ from unittest.mock import AsyncMock, Mock, patch
import pytest
from haystack import Document
from haystack.components.embedders import OpenAIDocumentEmbedder, SentenceTransformersDocumentEmbedder
from haystack.components.embedders import OpenAIDocumentEmbedder
from haystack.components.preprocessors import EmbeddingBasedDocumentSplitter
from haystack.utils import ComponentDevice
os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -359,18 +358,10 @@ class TestEmbeddingBasedDocumentSplitter:
assert result["init_parameters"]["max_length"] == 1000
assert "document_embedder" in result["init_parameters"]
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
@pytest.mark.slow
def test_split_document_with_multiple_topics(self, del_hf_env_vars, monkeypatch):
import torch
# Force CPU usage to avoid MPS memory issues
monkeypatch.setenv("PYTORCH_ENABLE_MPS_FALLBACK", "1")
torch.backends.mps.is_available = lambda: False
embedder = SentenceTransformersDocumentEmbedder(
model="sentence-transformers/all-MiniLM-L6-v2", device=ComponentDevice.from_str("cpu")
)
def test_split_document_with_multiple_topics(self):
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
splitter = EmbeddingBasedDocumentSplitter(
document_embedder=embedder, sentences_per_group=2, percentile=0.9, min_length=30, max_length=300
@@ -438,10 +429,10 @@ class TestEmbeddingBasedDocumentSplitter:
original = text
assert combined in original or original in combined
@pytest.mark.slow
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
def test_trailing_whitespace_is_preserved(self, del_hf_env_vars):
embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
def test_trailing_whitespace_is_preserved(self):
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
splitter = EmbeddingBasedDocumentSplitter(document_embedder=embedder, sentences_per_group=1)
@@ -482,10 +473,10 @@ class TestEmbeddingBasedDocumentSplitter:
result = await splitter.run_async(documents=[Document(content=text)])
assert result["documents"][0].content == text
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
@pytest.mark.slow
def test_no_extra_whitespaces_between_sentences(self, del_hf_env_vars):
embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
def test_no_extra_whitespaces_between_sentences(self):
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
splitter = EmbeddingBasedDocumentSplitter(
document_embedder=embedder, sentences_per_group=1, percentile=0.9, min_length=10, max_length=500
@@ -541,15 +532,15 @@ class TestEmbeddingBasedDocumentSplitter:
== "Machine learning has revolutionized many industries. Neural networks can process vast amounts of data. Deep learning models achieve remarkable accuracy on complex tasks." # noqa: E501
) # noqa: E501
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
@pytest.mark.slow
def test_split_large_splits_recursion(self, del_hf_env_vars):
def test_split_large_splits_recursion(self):
"""
Test that _split_large_splits() works correctly without infinite loops.
This test uses a longer text that will trigger the recursive splitting logic.
If the chunk cannot be split further, it is allowed to be larger than max_length.
"""
embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2", batch_size=32)
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
semantic_chunker = EmbeddingBasedDocumentSplitter(
document_embedder=embedder, sentences_per_group=5, percentile=0.95, min_length=50, max_length=1000
)
@@ -622,14 +613,14 @@ The history of software is closely tied to the development of digital computers
assert split_doc.meta["split_id"] == i
assert "page_number" in split_doc.meta
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
@pytest.mark.slow
def test_split_large_splits_actually_splits(self, del_hf_env_vars):
def test_split_large_splits_actually_splits(self):
"""
Test that _split_large_splits() actually works and can split long texts into multiple chunks.
This test uses a very long text that should be split into multiple chunks.
"""
embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2", batch_size=32)
embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
semantic_chunker = EmbeddingBasedDocumentSplitter(
document_embedder=embedder,
sentences_per_group=3,
test/components/rankers/test_sentence_transformers_diversity.py
-778
View File
@@ -1,778 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from unittest.mock import MagicMock, call, patch
import pytest
import torch
from haystack import Document, Pipeline
from haystack.components.rankers import SentenceTransformersDiversityRanker
from haystack.components.rankers.sentence_transformers_diversity import DiversityRankingSimilarity
from haystack.utils import ComponentDevice
from haystack.utils.auth import Secret
def mock_encode_response(texts, **kwargs):
if texts == ["city"]:
return torch.tensor([[1.0, 1.0]])
if texts == ["Eiffel Tower", "Berlin", "Bananas"]:
return torch.tensor([[1.0, 0.0], [0.8, 0.8], [0.0, 1.0]])
return torch.tensor([[0.0, 1.0]] * len(texts))
class TestSentenceTransformersDiversityRanker:
def test_init(self):
component = SentenceTransformersDiversityRanker()
assert component.model_name_or_path == "sentence-transformers/all-MiniLM-L6-v2"
assert component.top_k == 10
assert component.device == ComponentDevice.resolve_device(None)
assert component.similarity == DiversityRankingSimilarity.COSINE
assert component.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert component.query_prefix == ""
assert component.document_prefix == ""
assert component.query_suffix == ""
assert component.document_suffix == ""
assert component.meta_fields_to_embed == []
assert component.embedding_separator == "\n"
assert component.lambda_threshold == 0.5
def test_init_with_custom_parameters(self):
component = SentenceTransformersDiversityRanker(
model="sentence-transformers/msmarco-distilbert-base-v4",
top_k=5,
device=ComponentDevice.from_str("cuda:0"),
token=Secret.from_token("fake-api-token"),
similarity="dot_product",
query_prefix="query:",
document_prefix="document:",
query_suffix="query suffix",
document_suffix="document suffix",
meta_fields_to_embed=["meta_field"],
embedding_separator="--",
lambda_threshold=0,
)
assert component.model_name_or_path == "sentence-transformers/msmarco-distilbert-base-v4"
assert component.top_k == 5
assert component.device == ComponentDevice.from_str("cuda:0")
assert component.similarity == DiversityRankingSimilarity.DOT_PRODUCT
assert component.token == Secret.from_token("fake-api-token")
assert component.query_prefix == "query:"
assert component.document_prefix == "document:"
assert component.query_suffix == "query suffix"
assert component.document_suffix == "document suffix"
assert component.meta_fields_to_embed == ["meta_field"]
assert component.embedding_separator == "--"
assert component.lambda_threshold == 0
def test_to_dict(self):
component = SentenceTransformersDiversityRanker()
data = component.to_dict()
assert (
data["type"]
== "haystack.components.rankers.sentence_transformers_diversity.SentenceTransformersDiversityRanker"
)
assert data["init_parameters"]["model"] == "sentence-transformers/all-MiniLM-L6-v2"
assert data["init_parameters"]["top_k"] == 10
assert data["init_parameters"]["device"] == ComponentDevice.resolve_device(None).to_dict()
assert data["init_parameters"]["similarity"] == "cosine"
assert data["init_parameters"]["token"] == {
"env_vars": ["HF_API_TOKEN", "HF_TOKEN"],
"strict": False,
"type": "env_var",
}
assert data["init_parameters"]["query_prefix"] == ""
assert data["init_parameters"]["document_prefix"] == ""
assert data["init_parameters"]["query_suffix"] == ""
assert data["init_parameters"]["document_suffix"] == ""
assert data["init_parameters"]["meta_fields_to_embed"] == []
assert data["init_parameters"]["embedding_separator"] == "\n"
assert data["init_parameters"]["strategy"] == "greedy_diversity_order"
def test_from_dict(self):
data = {
"type": "haystack.components.rankers.sentence_transformers_diversity.SentenceTransformersDiversityRanker",
"init_parameters": {
"model": "sentence-transformers/all-MiniLM-L6-v2",
"top_k": 10,
"device": ComponentDevice.resolve_device(None).to_dict(),
"similarity": "cosine",
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"query_prefix": "",
"document_prefix": "",
"query_suffix": "",
"document_suffix": "",
"meta_fields_to_embed": [],
"embedding_separator": "\n",
},
}
ranker = SentenceTransformersDiversityRanker.from_dict(data)
assert ranker.model_name_or_path == "sentence-transformers/all-MiniLM-L6-v2"
assert ranker.top_k == 10
assert ranker.device == ComponentDevice.resolve_device(None)
assert ranker.similarity == DiversityRankingSimilarity.COSINE
assert ranker.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert ranker.query_prefix == ""
assert ranker.document_prefix == ""
assert ranker.query_suffix == ""
assert ranker.document_suffix == ""
assert ranker.meta_fields_to_embed == []
assert ranker.embedding_separator == "\n"
def test_from_dict_none_device(self):
data = {
"type": "haystack.components.rankers.sentence_transformers_diversity.SentenceTransformersDiversityRanker",
"init_parameters": {
"model": "sentence-transformers/all-MiniLM-L6-v2",
"top_k": 10,
"device": None,
"similarity": "cosine",
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"query_prefix": "",
"document_prefix": "",
"query_suffix": "",
"document_suffix": "",
"meta_fields_to_embed": [],
"embedding_separator": "\n",
},
}
ranker = SentenceTransformersDiversityRanker.from_dict(data)
assert ranker.model_name_or_path == "sentence-transformers/all-MiniLM-L6-v2"
assert ranker.top_k == 10
assert ranker.device == ComponentDevice.resolve_device(None)
assert ranker.similarity == DiversityRankingSimilarity.COSINE
assert ranker.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert ranker.query_prefix == ""
assert ranker.document_prefix == ""
assert ranker.query_suffix == ""
assert ranker.document_suffix == ""
assert ranker.meta_fields_to_embed == []
assert ranker.embedding_separator == "\n"
def test_from_dict_no_default_parameters(self):
data = {
"type": "haystack.components.rankers.sentence_transformers_diversity.SentenceTransformersDiversityRanker",
"init_parameters": {},
}
ranker = SentenceTransformersDiversityRanker.from_dict(data)
assert ranker.model_name_or_path == "sentence-transformers/all-MiniLM-L6-v2"
assert ranker.top_k == 10
assert ranker.device == ComponentDevice.resolve_device(None)
assert ranker.similarity == DiversityRankingSimilarity.COSINE
assert ranker.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert ranker.query_prefix == ""
assert ranker.document_prefix == ""
assert ranker.query_suffix == ""
assert ranker.document_suffix == ""
assert ranker.meta_fields_to_embed == []
assert ranker.embedding_separator == "\n"
def test_to_dict_with_custom_parameters(self):
component = SentenceTransformersDiversityRanker(
model="sentence-transformers/msmarco-distilbert-base-v4",
top_k=5,
device=ComponentDevice.from_str("cuda:0"),
token=Secret.from_env_var("ENV_VAR", strict=False),
similarity="dot_product",
query_prefix="query:",
document_prefix="document:",
query_suffix="query suffix",
document_suffix="document suffix",
meta_fields_to_embed=["meta_field"],
embedding_separator="--",
)
data = component.to_dict()
assert (
data["type"]
== "haystack.components.rankers.sentence_transformers_diversity.SentenceTransformersDiversityRanker"
)
assert data["init_parameters"]["model"] == "sentence-transformers/msmarco-distilbert-base-v4"
assert data["init_parameters"]["top_k"] == 5
assert data["init_parameters"]["device"] == ComponentDevice.from_str("cuda:0").to_dict()
assert data["init_parameters"]["token"] == {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"}
assert data["init_parameters"]["similarity"] == "dot_product"
assert data["init_parameters"]["query_prefix"] == "query:"
assert data["init_parameters"]["document_prefix"] == "document:"
assert data["init_parameters"]["query_suffix"] == "query suffix"
assert data["init_parameters"]["document_suffix"] == "document suffix"
assert data["init_parameters"]["meta_fields_to_embed"] == ["meta_field"]
assert data["init_parameters"]["embedding_separator"] == "--"
assert data["init_parameters"]["strategy"] == "greedy_diversity_order"
def test_from_dict_with_custom_init_parameters(self):
data = {
"type": "haystack.components.rankers.sentence_transformers_diversity.SentenceTransformersDiversityRanker",
"init_parameters": {
"model": "sentence-transformers/msmarco-distilbert-base-v4",
"top_k": 5,
"device": ComponentDevice.from_str("cuda:0").to_dict(),
"token": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
"similarity": "dot_product",
"query_prefix": "query:",
"document_prefix": "document:",
"query_suffix": "query suffix",
"document_suffix": "document suffix",
"meta_fields_to_embed": ["meta_field"],
"embedding_separator": "--",
},
}
ranker = SentenceTransformersDiversityRanker.from_dict(data)
assert ranker.model_name_or_path == "sentence-transformers/msmarco-distilbert-base-v4"
assert ranker.top_k == 5
assert ranker.device == ComponentDevice.from_str("cuda:0")
assert ranker.similarity == DiversityRankingSimilarity.DOT_PRODUCT
assert ranker.token == Secret.from_env_var("ENV_VAR", strict=False)
assert ranker.query_prefix == "query:"
assert ranker.document_prefix == "document:"
assert ranker.query_suffix == "query suffix"
assert ranker.document_suffix == "document suffix"
assert ranker.meta_fields_to_embed == ["meta_field"]
assert ranker.embedding_separator == "--"
def test_run_invalid_similarity(self):
"""
Tests that run method raises ValueError if similarity is incorrect
"""
similarity = "incorrect"
with pytest.raises(ValueError, match="Unknown similarity metric"):
SentenceTransformersDiversityRanker(model="sentence-transformers/all-MiniLM-L6-v2", similarity=similarity)
def test_run_invalid_strategy(self):
"""
Tests that run method raises ValueError if strategy is incorrect
"""
strategy = "incorrect"
with pytest.raises(ValueError, match="Unknown strategy"):
SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity="cosine", strategy=strategy
)
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_warm_up(self, similarity, del_hf_env_vars):
"""
Test that ranker loads the SentenceTransformer model correctly during warm up.
"""
mock_model_class = MagicMock()
mock_model_instance = MagicMock()
mock_model_class.return_value = mock_model_instance
with patch(
"haystack.components.rankers.sentence_transformers_diversity.SentenceTransformer", new=mock_model_class
):
ranker = SentenceTransformersDiversityRanker(model="mock_model_name", similarity=similarity)
assert ranker.model is None
ranker.warm_up()
mock_model_class.assert_called_once_with(
model_name_or_path="mock_model_name",
device=ComponentDevice.resolve_device(None).to_torch_str(),
token=None,
model_kwargs=None,
tokenizer_kwargs=None,
config_kwargs=None,
backend="torch",
)
assert ranker.model == mock_model_instance
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_empty_query(self, similarity):
"""
Test that ranker can be run with an empty query.
"""
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", top_k=3, similarity=similarity
)
ranker.model = MagicMock()
ranker.model.encode = MagicMock(side_effect=mock_encode_response)
documents = [Document(content="doc1"), Document(content="doc2")]
result = ranker.run(query="", documents=documents)
ranked_docs = result["documents"]
assert isinstance(ranked_docs, list)
assert len(ranked_docs) == 2
assert all(isinstance(doc, Document) for doc in ranked_docs)
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_top_k(self, similarity):
"""
Test that run method returns the correct number of documents for different top_k values passed at
initialization and runtime.
"""
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity=similarity, top_k=3
)
ranker.model = MagicMock()
ranker.model.encode = MagicMock(side_effect=mock_encode_response)
query = "test query"
documents = [
Document(content="doc1"),
Document(content="doc2"),
Document(content="doc3"),
Document(content="doc4"),
]
result = ranker.run(query=query, documents=documents)
ranked_docs = result["documents"]
assert isinstance(ranked_docs, list)
assert len(ranked_docs) == 3
assert all(isinstance(doc, Document) for doc in ranked_docs)
# Passing a different top_k at runtime
result = ranker.run(query=query, documents=documents, top_k=2)
ranked_docs = result["documents"]
assert isinstance(ranked_docs, list)
assert len(ranked_docs) == 2
assert all(isinstance(doc, Document) for doc in ranked_docs)
def test_run_deduplicates_documents(self):
ranker = SentenceTransformersDiversityRanker()
ranker.model = MagicMock()
ranker.model.encode = MagicMock(side_effect=mock_encode_response)
documents = [
Document(id="duplicate", content="keep me", score=0.9),
Document(id="duplicate", content="drop me", score=0.1),
Document(id="unique", content="unique"),
]
result = ranker.run(query="test", documents=documents)
assert len(result["documents"]) == 2
assert result["documents"][0].content == "keep me"
assert result["documents"][1].content == "unique"
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_negative_top_k_at_init(self, similarity):
"""
Tests that run method raises an error for negative top-k set at init.
"""
with pytest.raises(ValueError, match="top_k must be > 0, but got"):
SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity=similarity, top_k=-5
)
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_top_k_is_none_at_init(self, similarity):
"""
Tests that run method raises an error for top-k set to None at init.
"""
with pytest.raises(ValueError, match="top_k must be > 0, but got"):
SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity=similarity, top_k=None
)
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_negative_top_k(self, similarity):
"""
Tests that run method raises an error for negative top-k set at runtime.
"""
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity=similarity, top_k=10
)
ranker.model = MagicMock()
query = "test"
documents = [Document(content="doc1"), Document(content="doc2"), Document(content="doc3")]
with pytest.raises(ValueError, match="top_k must be > 0"):
ranker.run(query=query, documents=documents, top_k=-5)
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_top_k_is_none(self, similarity):
"""
Tests that run method returns the correct order of documents for top-k set to None.
"""
# Setting top_k to None is ignored during runtime, it should use top_k set at init.
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity=similarity, top_k=2
)
ranker.model = MagicMock()
ranker.model.encode = MagicMock(side_effect=mock_encode_response)
query = "test"
documents = [Document(content="doc1"), Document(content="doc2"), Document(content="doc3")]
result = ranker.run(query=query, documents=documents, top_k=None)
assert len(result["documents"]) == 2
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_no_documents_provided(self, similarity):
"""
Test that run method returns an empty list if no documents are supplied.
"""
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity=similarity
)
ranker.model = MagicMock()
query = "test query"
documents = []
results = ranker.run(query=query, documents=documents)
assert len(results["documents"]) == 0
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_with_less_documents_than_top_k(self, similarity):
"""
Tests that run method returns the correct number of documents for top_k values greater than number of documents.
"""
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity=similarity, top_k=5
)
ranker.model = MagicMock()
ranker.model.encode = MagicMock(side_effect=mock_encode_response)
query = "test"
documents = [Document(content="doc1"), Document(content="doc2"), Document(content="doc3")]
result = ranker.run(query=query, documents=documents)
assert len(result["documents"]) == 3
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_single_document_corner_case(self, similarity):
"""
Tests that run method returns the correct number of documents for a single document
"""
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity=similarity
)
ranker.model = MagicMock()
ranker.model.encode = MagicMock(side_effect=mock_encode_response)
query = "test"
documents = [Document(content="doc1")]
result = ranker.run(query=query, documents=documents)
assert len(result["documents"]) == 1
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_prepare_texts_to_embed(self, similarity):
"""
Test creation of texts to embed from documents with meta fields, document prefix and suffix.
"""
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2",
similarity=similarity,
document_prefix="test doc: ",
document_suffix=" end doc.",
meta_fields_to_embed=["meta_field"],
embedding_separator="\n",
)
documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)]
texts = ranker._prepare_texts_to_embed(documents=documents)
assert texts == [
"test doc: meta_value 0\ndocument number 0 end doc.",
"test doc: meta_value 1\ndocument number 1 end doc.",
"test doc: meta_value 2\ndocument number 2 end doc.",
"test doc: meta_value 3\ndocument number 3 end doc.",
"test doc: meta_value 4\ndocument number 4 end doc.",
]
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_prepare_texts_to_embed_preserves_falsy_meta_values(self, similarity):
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2",
similarity=similarity,
document_prefix="test doc: ",
document_suffix=" end doc.",
meta_fields_to_embed=["rating", "is_available", "missing", "none_value"],
embedding_separator="\n",
)
documents = [Document(content="document", meta={"rating": 0, "is_available": False, "none_value": None})]
texts = ranker._prepare_texts_to_embed(documents=documents)
assert texts == ["test doc: 0\nFalse\ndocument end doc."]
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_encode_text(self, similarity):
"""
Test addition of suffix and prefix to the query and documents when creating embeddings.
"""
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2",
similarity=similarity,
query_prefix="test query: ",
query_suffix=" end query.",
document_prefix="test doc: ",
document_suffix=" end doc.",
meta_fields_to_embed=["meta_field"],
embedding_separator="\n",
)
query = "query"
documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)]
ranker.model = MagicMock()
ranker.model.encode = MagicMock(side_effect=mock_encode_response)
ranker.run(query=query, documents=documents)
assert ranker.model.encode.call_count == 2
ranker.model.assert_has_calls(
[
call.encode(
[
"test doc: meta_value 0\ndocument number 0 end doc.",
"test doc: meta_value 1\ndocument number 1 end doc.",
"test doc: meta_value 2\ndocument number 2 end doc.",
"test doc: meta_value 3\ndocument number 3 end doc.",
"test doc: meta_value 4\ndocument number 4 end doc.",
],
convert_to_tensor=True,
),
call.encode(["test query: query end query."], convert_to_tensor=True),
]
)
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_greedy_diversity_order(self, similarity):
"""
Tests that the given list of documents is ordered to maximize diversity.
"""
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity=similarity
)
query = "city"
documents = [Document(content="Eiffel Tower"), Document(content="Berlin"), Document(content="Bananas")]
ranker.model = MagicMock()
ranker.model.encode = MagicMock(side_effect=mock_encode_response)
ranked_docs = ranker._greedy_diversity_order(query=query, documents=documents)
ranked_text = " ".join([doc.content for doc in ranked_docs])
assert ranked_text == "Berlin Eiffel Tower Bananas"
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_maximum_margin_relevance(self, similarity):
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity=similarity
)
ranker.model = MagicMock()
ranker.model.encode = MagicMock(side_effect=mock_encode_response)
query = "city"
documents = [Document(content="Eiffel Tower"), Document(content="Berlin"), Document(content="Bananas")]
ranker.model = MagicMock()
ranker.model.encode = MagicMock(side_effect=mock_encode_response)
ranked_docs = ranker._maximum_margin_relevance(query=query, documents=documents, lambda_threshold=0, top_k=3)
ranked_text = " ".join([doc.content for doc in ranked_docs])
assert ranked_text == "Berlin Eiffel Tower Bananas"
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_maximum_margin_relevance_with_given_lambda_threshold(self, similarity):
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity=similarity
)
ranker.model = MagicMock()
ranker.model.encode = MagicMock(side_effect=mock_encode_response)
query = "city"
documents = [Document(content="Eiffel Tower"), Document(content="Berlin"), Document(content="Bananas")]
ranker.model = MagicMock()
ranker.model.encode = MagicMock(side_effect=mock_encode_response)
ranked_docs = ranker._maximum_margin_relevance(query=query, documents=documents, lambda_threshold=1, top_k=3)
ranked_text = " ".join([doc.content for doc in ranked_docs])
assert ranked_text == "Berlin Eiffel Tower Bananas"
def test_pipeline_serialise_deserialise(self):
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2", similarity="cosine", top_k=5
)
pipe = Pipeline()
pipe.add_component("ranker", ranker)
pipe_serialized = pipe.dumps()
assert Pipeline.loads(pipe_serialized) == pipe
@pytest.mark.integration
@pytest.mark.slow
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_real_world_use_case(self, similarity, del_hf_env_vars):
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers-testing/stsb-bert-tiny-safetensors", similarity=similarity
)
query = "What are the reasons for long-standing animosities between Russia and Poland?"
doc1 = Document(
"One of the earliest known events in Russian-Polish history dates back to 981, when the Grand Prince of "
"Kiev , Vladimir Svyatoslavich , seized the Cherven Cities from the Duchy of Poland . The relationship "
"between two by that time was mostly close and cordial, as there had been no serious wars between both. "
"In 966, Poland accepted Christianity from Rome while Kievan Rus' —the ancestor of Russia, Ukraine and "
"Belarus—was Christianized by Constantinople. In 1054, the internal Christian divide formally split the "
"Church into the Catholic and Orthodox branches separating the Poles from the Eastern Slavs."
)
doc2 = Document(
"Since the fall of the Soviet Union , with Lithuania , Ukraine and Belarus regaining independence, the "
"Polish Russian border has mostly been replaced by borders with the respective countries, but there still "
"is a 210 km long border between Poland and the Kaliningrad Oblast"
)
doc3 = Document(
"As part of Poland's plans to become fully energy independent from Russia within the next years, Piotr "
"Wozniak, president of state-controlled oil and gas company PGNiG , stated in February 2019: 'The strategy "
"of the company is just to forget about Eastern suppliers and especially about Gazprom .'[53] In 2020, the "
"Stockholm Arbitrary Tribunal ruled that PGNiG's long-term contract gas price with Gazprom linked to oil "
"prices should be changed to approximate the Western European gas market price, backdated to 1 November "
"2014 when PGNiG requested a price review under the contract. Gazprom had to refund about $1.5 billion to "
"PGNiG."
)
doc4 = Document(
"Both Poland and Russia had accused each other for their historical revisionism . Russia has repeatedly "
"accused Poland for not honoring Soviet Red Army soldiers fallen in World War II for Poland, notably in "
"2017, in which Poland was thought on 'attempting to impose its own version of history' after Moscow was "
"not allowed to join an international effort to renovate a World War II museum at Sobibór , site of a "
"notorious Sobibor extermination camp."
)
doc5 = Document(
"President of Russia Vladimir Putin and Prime Minister of Poland Leszek Miller in 2002 Modern Polish "
"Russian relations begin with the fall of communism in1989 in Poland ( Solidarity and the Polish Round "
"Table Agreement ) and 1991 in Russia ( dissolution of the Soviet Union ). With a new democratic "
"government after the 1989 elections , Poland regained full sovereignty, [2] and what was the Soviet "
"Union, became 15 newly independent states , including the Russian Federation . Relations between modern "
"Poland and Russia suffer from constant ups and downs."
)
doc6 = Document(
"Soviet influence in Poland finally ended with the Round Table Agreement of 1989 guaranteeing free "
"elections in Poland, the Revolutions of 1989 against Soviet-sponsored Communist governments in the "
"Eastern Block , and finally the formal dissolution of the Warsaw Pact."
)
doc7 = Document(
"Dmitry Medvedev and then Polish Prime Minister Donald Tusk , 6 December 2010 BBC News reported that one "
"of the main effects of the 2010 Polish Air Force Tu-154 crash would be the impact it has on "
"Russian-Polish relations. [38] It was thought if the inquiry into the crash were not transparent, it "
"would increase suspicions toward Russia in Poland."
)
doc8 = Document(
"Soviet control over the Polish People's Republic lessened after Stalin's death and Gomulka's Thaw , and "
"ceased completely after the fall of the communist government in Poland in late 1989, although the "
"Soviet-Russian Northern Group of Forces did not leave Polish soil until 1993. The continuing Soviet "
"military presence allowed the Soviet Union to heavily influence Polish politics."
)
documents = [doc1, doc2, doc3, doc4, doc5, doc6, doc7, doc8]
result = ranker.run(query=query, documents=documents)
expected_order = [doc5, doc7, doc3, doc1, doc4, doc2, doc6, doc8]
expected_content = " ".join([doc.content or "" for doc in expected_order])
result_content = " ".join([doc.content or "" for doc in result["documents"]])
# Check the order of ranked documents by comparing the content of the ranked documents
assert result_content == expected_content
@pytest.mark.integration
@pytest.mark.slow
@pytest.mark.parametrize("similarity", ["dot_product", "cosine"])
def test_run_with_maximum_margin_relevance_strategy(self, similarity, del_hf_env_vars):
query = "renewable energy sources"
docs = [
Document(content="18th-century French literature"),
Document(content="Solar power generation"),
Document(content="Ancient Egyptian hieroglyphics"),
Document(content="Wind turbine technology"),
Document(content="Baking sourdough bread"),
Document(content="Hydroelectric dam systems"),
Document(content="Geothermal energy extraction"),
Document(content="Biomass fuel production"),
]
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers-testing/stsb-bert-tiny-safetensors",
similarity=similarity,
strategy="maximum_margin_relevance",
)
# lambda_threshold=1, the most relevant document should be returned first
results = ranker.run(query=query, documents=docs, lambda_threshold=1, top_k=len(docs))
expected = [
"Geothermal energy extraction",
"Wind turbine technology",
"Solar power generation",
"Biomass fuel production",
"Hydroelectric dam systems",
"Baking sourdough bread",
"Ancient Egyptian hieroglyphics",
"18th-century French literature",
]
assert [doc.content for doc in results["documents"]] == expected
# lambda_threshold=0, after the most relevant one, diverse documents should be returned
results = ranker.run(query=query, documents=docs, lambda_threshold=0, top_k=len(docs))
expected = [
"Geothermal energy extraction",
"18th-century French literature",
"Ancient Egyptian hieroglyphics",
"Baking sourdough bread",
"Solar power generation",
"Biomass fuel production",
"Hydroelectric dam systems",
"Wind turbine technology",
]
assert [doc.content for doc in results["documents"]] == expected
@patch("haystack.components.rankers.sentence_transformers_diversity.SentenceTransformer")
def test_model_onnx_backend(self, mocked_sentence_transformer):
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cpu"),
model_kwargs={"file_name": "onnx/model.onnx"},
backend="onnx",
)
ranker.warm_up()
mocked_sentence_transformer.assert_called_once_with(
model_name_or_path="sentence-transformers/all-MiniLM-L6-v2",
device="cpu",
token=None,
model_kwargs={"file_name": "onnx/model.onnx"},
tokenizer_kwargs=None,
config_kwargs=None,
backend="onnx",
)
@patch("haystack.components.rankers.sentence_transformers_diversity.SentenceTransformer")
def test_model_openvino_backend(self, mocked_sentence_transformer):
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cpu"),
model_kwargs={"file_name": "openvino/openvino_model.xml"},
backend="openvino",
)
ranker.warm_up()
mocked_sentence_transformer.assert_called_once_with(
model_name_or_path="sentence-transformers/all-MiniLM-L6-v2",
device="cpu",
token=None,
model_kwargs={"file_name": "openvino/openvino_model.xml"},
tokenizer_kwargs=None,
config_kwargs=None,
backend="openvino",
)
@patch("haystack.components.rankers.sentence_transformers_diversity.SentenceTransformer")
@pytest.mark.parametrize("model_kwargs", [{"torch_dtype": "float16"}, {"torch_dtype": "bfloat16"}])
def test_dtype_on_gpu(self, mocked_sentence_transformer, model_kwargs):
ranker = SentenceTransformersDiversityRanker(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cuda:0"),
model_kwargs=model_kwargs,
)
ranker.warm_up()
mocked_sentence_transformer.assert_called_once_with(
model_name_or_path="sentence-transformers/all-MiniLM-L6-v2",
device="cuda:0",
token=None,
model_kwargs=model_kwargs,
tokenizer_kwargs=None,
config_kwargs=None,
backend="torch",
)
test/components/rankers/test_sentence_transformers_similarity.py
-490
View File
@@ -1,490 +0,0 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
import torch
from haystack import Document
from haystack.components.rankers.sentence_transformers_similarity import SentenceTransformersSimilarityRanker
from haystack.utils.auth import Secret
from haystack.utils.device import ComponentDevice
class TestSentenceTransformersSimilarityRanker:
def test_init_invalid_top_k(self):
with pytest.raises(ValueError):
SentenceTransformersSimilarityRanker(top_k=-1)
@patch("haystack.components.rankers.sentence_transformers_similarity.CrossEncoder")
def test_init_warm_up_torch_backend(self, mocked_cross_encoder):
ranker = SentenceTransformersSimilarityRanker(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cpu"),
backend="torch",
trust_remote_code=True,
)
ranker.warm_up()
mocked_cross_encoder.assert_called_once_with(
model_name_or_path="sentence-transformers/all-MiniLM-L6-v2",
device="cpu",
token=None,
trust_remote_code=True,
model_kwargs=None,
tokenizer_kwargs=None,
config_kwargs=None,
backend="torch",
)
@patch("haystack.components.rankers.sentence_transformers_similarity.CrossEncoder")
def test_init_warm_up_onnx_backend(self, mocked_cross_encoder):
onnx_ranker = SentenceTransformersSimilarityRanker(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cpu"),
backend="onnx",
)
onnx_ranker.warm_up()
mocked_cross_encoder.assert_called_once_with(
model_name_or_path="sentence-transformers/all-MiniLM-L6-v2",
device="cpu",
token=None,
trust_remote_code=False,
model_kwargs=None,
tokenizer_kwargs=None,
config_kwargs=None,
backend="onnx",
)
@patch("haystack.components.rankers.sentence_transformers_similarity.CrossEncoder")
def test_init_warm_up_openvino_backend(self, mocked_cross_encoder):
openvino_ranker = SentenceTransformersSimilarityRanker(
model="sentence-transformers/all-MiniLM-L6-v2",
token=None,
device=ComponentDevice.from_str("cpu"),
backend="openvino",
)
openvino_ranker.warm_up()
mocked_cross_encoder.assert_called_once_with(
model_name_or_path="sentence-transformers/all-MiniLM-L6-v2",
device="cpu",
token=None,
trust_remote_code=False,
model_kwargs=None,
tokenizer_kwargs=None,
config_kwargs=None,
backend="openvino",
)
def test_to_dict(self):
component = SentenceTransformersSimilarityRanker()
data = component.to_dict()
assert data == {
"type": "haystack.components.rankers.sentence_transformers_similarity.SentenceTransformersSimilarityRanker",
"init_parameters": {
"device": ComponentDevice.resolve_device(None).to_dict(),
"top_k": 10,
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"query_prefix": "",
"query_suffix": "",
"document_prefix": "",
"document_suffix": "",
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2",
"meta_fields_to_embed": [],
"embedding_separator": "\n",
"scale_score": True,
"score_threshold": None,
"trust_remote_code": False,
"model_kwargs": None,
"tokenizer_kwargs": None,
"config_kwargs": None,
"backend": "torch",
"batch_size": 16,
},
}
def test_to_dict_with_custom_init_parameters(self):
component = SentenceTransformersSimilarityRanker(
model="my_model",
device=ComponentDevice.from_str("cuda:0"),
token=Secret.from_env_var("ENV_VAR", strict=False),
top_k=5,
query_prefix="query_instruction: ",
query_suffix="\n",
document_prefix="document_instruction: ",
document_suffix="\n",
scale_score=False,
score_threshold=0.01,
trust_remote_code=True,
model_kwargs={"torch_dtype": torch.float16},
tokenizer_kwargs={"model_max_length": 512},
batch_size=32,
)
data = component.to_dict()
assert data == {
"type": "haystack.components.rankers.sentence_transformers_similarity.SentenceTransformersSimilarityRanker",
"init_parameters": {
"device": {"type": "single", "device": "cuda:0"},
"model": "my_model",
"token": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
"top_k": 5,
"query_prefix": "query_instruction: ",
"query_suffix": "\n",
"document_prefix": "document_instruction: ",
"document_suffix": "\n",
"meta_fields_to_embed": [],
"embedding_separator": "\n",
"scale_score": False,
"score_threshold": 0.01,
"trust_remote_code": True,
"model_kwargs": {"torch_dtype": "torch.float16"},
"tokenizer_kwargs": {"model_max_length": 512},
"config_kwargs": None,
"backend": "torch",
"batch_size": 32,
},
}
def test_to_dict_with_quantization_options(self):
component = SentenceTransformersSimilarityRanker(
model_kwargs={
"load_in_4bit": True,
"bnb_4bit_use_double_quant": True,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_compute_dtype": torch.bfloat16,
}
)
data = component.to_dict()
assert data == {
"type": "haystack.components.rankers.sentence_transformers_similarity.SentenceTransformersSimilarityRanker",
"init_parameters": {
"device": ComponentDevice.resolve_device(None).to_dict(),
"top_k": 10,
"query_prefix": "",
"query_suffix": "",
"document_prefix": "",
"document_suffix": "",
"token": {"env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False, "type": "env_var"},
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2",
"meta_fields_to_embed": [],
"embedding_separator": "\n",
"scale_score": True,
"score_threshold": None,
"trust_remote_code": False,
"model_kwargs": {
"load_in_4bit": True,
"bnb_4bit_use_double_quant": True,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_compute_dtype": "torch.bfloat16",
},
"tokenizer_kwargs": None,
"config_kwargs": None,
"backend": "torch",
"batch_size": 16,
},
}
def test_from_dict(self):
data = {
"type": "haystack.components.rankers.sentence_transformers_similarity.SentenceTransformersSimilarityRanker",
"init_parameters": {
"device": ComponentDevice.resolve_device(None).to_dict(),
"model": "my_model",
"token": None,
"top_k": 5,
"query_prefix": "",
"document_prefix": "",
"meta_fields_to_embed": [],
"embedding_separator": "\n",
"scale_score": False,
"score_threshold": 0.01,
"trust_remote_code": False,
"model_kwargs": {"torch_dtype": "torch.float16"},
"tokenizer_kwargs": None,
"config_kwargs": None,
"backend": "torch",
"batch_size": 32,
},
}
component = SentenceTransformersSimilarityRanker.from_dict(data)
assert component.device == ComponentDevice.resolve_device(None)
assert component.model == "my_model"
assert component.token is None
assert component.top_k == 5
assert component.query_prefix == ""
assert component.document_prefix == ""
assert component.meta_fields_to_embed == []
assert component.embedding_separator == "\n"
assert not component.scale_score
assert component.score_threshold == 0.01
assert component.trust_remote_code is False
assert component.model_kwargs == {"torch_dtype": torch.float16}
assert component.tokenizer_kwargs is None
assert component.config_kwargs is None
assert component.batch_size == 32
def test_from_dict_no_default_parameters(self):
data = {
"type": "haystack.components.rankers.sentence_transformers_similarity.SentenceTransformersSimilarityRanker",
"init_parameters": {},
}
component = SentenceTransformersSimilarityRanker.from_dict(data)
assert component.device == ComponentDevice.resolve_device(None)
assert component.model == "cross-encoder/ms-marco-MiniLM-L-6-v2"
assert component.token == Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False)
assert component.top_k == 10
assert component.query_prefix == ""
assert component.document_prefix == ""
assert component.meta_fields_to_embed == []
assert component.embedding_separator == "\n"
assert component.scale_score
assert component.score_threshold is None
assert component.trust_remote_code is False
assert component.model_kwargs is None
assert component.tokenizer_kwargs is None
assert component.config_kwargs is None
assert component.backend == "torch"
assert component.batch_size == 16
@patch("haystack.components.rankers.sentence_transformers_similarity.CrossEncoder")
def test_warm_up(self, mock_cross_encoder):
ranker = SentenceTransformersSimilarityRanker()
mock_cross_encoder.assert_not_called()
ranker.warm_up()
mock_cross_encoder.assert_called()
def test_run_invalid_top_k(self):
ranker = SentenceTransformersSimilarityRanker()
ranker._cross_encoder = MagicMock()
with pytest.raises(ValueError):
ranker.run(query="test", documents=[Document(content="document")], top_k=-1)
def test_returns_empty_list_if_no_documents_are_provided(self):
ranker = SentenceTransformersSimilarityRanker()
ranker._cross_encoder = MagicMock()
output = ranker.run(query="City in Germany", documents=[])
assert not output["documents"]
def test_embed_meta(self):
ranker = SentenceTransformersSimilarityRanker(
model="model", meta_fields_to_embed=["meta_field"], embedding_separator="\n"
)
mock_cross_encoder = MagicMock()
ranker._cross_encoder = mock_cross_encoder
documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)]
ranker.run(query="test", documents=documents)
_, kwargs = mock_cross_encoder.rank.call_args
assert kwargs["query"] == "test"
assert kwargs["documents"] == [
"meta_value 0\ndocument number 0",
"meta_value 1\ndocument number 1",
"meta_value 2\ndocument number 2",
"meta_value 3\ndocument number 3",
"meta_value 4\ndocument number 4",
]
assert kwargs["batch_size"] == 16
assert isinstance(kwargs["activation_fn"], torch.nn.Sigmoid)
assert kwargs["convert_to_numpy"] is True
assert kwargs["return_documents"] is False
def test_embed_meta_preserves_falsy_values(self):
ranker = SentenceTransformersSimilarityRanker(
model="model",
meta_fields_to_embed=["rating", "is_available", "missing", "none_value"],
embedding_separator="\n",
)
mock_cross_encoder = MagicMock()
mock_cross_encoder.rank.return_value = [{"corpus_id": 0, "score": 1.0}]
ranker._cross_encoder = mock_cross_encoder
documents = [Document(content="document", meta={"rating": 0, "is_available": False, "none_value": None})]
ranker.run(query="test", documents=documents)
_, kwargs = mock_cross_encoder.rank.call_args
assert kwargs["documents"] == ["0\nFalse\ndocument"]
def test_prefix(self):
ranker = SentenceTransformersSimilarityRanker(
model="model", query_prefix="query_instruction: ", document_prefix="document_instruction: "
)
mock_cross_encoder = MagicMock()
ranker._cross_encoder = mock_cross_encoder
documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)]
ranker.run(query="test", documents=documents)
_, kwargs = mock_cross_encoder.rank.call_args
assert kwargs["query"] == "query_instruction: test"
assert kwargs["documents"] == [
"document_instruction: document number 0",
"document_instruction: document number 1",
"document_instruction: document number 2",
"document_instruction: document number 3",
"document_instruction: document number 4",
]
assert kwargs["batch_size"] == 16
assert isinstance(kwargs["activation_fn"], torch.nn.Sigmoid)
assert kwargs["convert_to_numpy"] is True
def test_suffix(self):
ranker = SentenceTransformersSimilarityRanker(
model="model", query_suffix="\nquery_suffix", document_suffix="\ndocument_suffix"
)
mock_cross_encoder = MagicMock()
ranker._cross_encoder = mock_cross_encoder
documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)]
ranker.run(query="test", documents=documents)
_, kwargs = mock_cross_encoder.rank.call_args
assert kwargs["query"] == "test\nquery_suffix"
assert kwargs["documents"] == [
"document number 0\ndocument_suffix",
"document number 1\ndocument_suffix",
"document number 2\ndocument_suffix",
"document number 3\ndocument_suffix",
"document number 4\ndocument_suffix",
]
assert kwargs["batch_size"] == 16
assert isinstance(kwargs["activation_fn"], torch.nn.Sigmoid)
assert kwargs["convert_to_numpy"] is True
def test_scale_score_false(self):
mock_cross_encoder = MagicMock()
ranker = SentenceTransformersSimilarityRanker(model="model", scale_score=False)
ranker._cross_encoder = mock_cross_encoder
mock_cross_encoder.rank.return_value = [{"score": -10.6859, "corpus_id": 0}, {"score": -8.9874, "corpus_id": 1}]
documents = [Document(content="document number 0"), Document(content="document number 1")]
out = ranker.run(query="test", documents=documents)
assert out["documents"][0].score == pytest.approx(-10.6859, abs=1e-4)
assert out["documents"][1].score == pytest.approx(-8.9874, abs=1e-4)
def test_score_threshold(self):
mock_cross_encoder = MagicMock()
ranker = SentenceTransformersSimilarityRanker(model="model", scale_score=False, score_threshold=0.1)
ranker._cross_encoder = mock_cross_encoder
mock_cross_encoder.rank.return_value = [{"score": 0.955, "corpus_id": 0}, {"score": 0.001, "corpus_id": 1}]
documents = [Document(content="document number 0"), Document(content="document number 1")]
out = ranker.run(query="test", documents=documents)
assert len(out["documents"]) == 1
documents = [Document(content="document number 0"), Document(content="document number 1")]
out = ranker.run(query="test", documents=documents)
assert len(out["documents"]) == 1
def test_scores_cast_to_python_float_when_numpy_scalars_returned(self):
mock_cross_encoder = MagicMock()
ranker = SentenceTransformersSimilarityRanker(model="model")
ranker._cross_encoder = mock_cross_encoder
# Simulate backend returning numpy scalar types
mock_cross_encoder.rank.return_value = [
{"score": np.float32(0.123), "corpus_id": 0},
{"score": np.float64(0.456), "corpus_id": 1},
]
documents = [Document(content="doc 0"), Document(content="doc 1")]
out = ranker.run(query="test", documents=documents)
assert len(out["documents"]) == 2
for d in out["documents"]:
assert isinstance(d.score, float)
def test_run_deduplicates_documents(self):
ranker = SentenceTransformersSimilarityRanker()
ranker._cross_encoder = MagicMock()
ranker._cross_encoder.rank.return_value = [{"score": 0.42, "corpus_id": 0}, {"score": 0.12, "corpus_id": 1}]
documents = [
Document(id="duplicate", content="keep me", score=0.9),
Document(id="duplicate", content="drop me", score=0.1),
Document(id="unique", content="unique"),
]
result = ranker.run(query="test", documents=documents)
assert len(result["documents"]) == 2
assert result["documents"][0].content == "keep me"
assert result["documents"][1].content == "unique"
@pytest.mark.integration
@pytest.mark.slow
def test_run(self, del_hf_env_vars):
ranker = SentenceTransformersSimilarityRanker(model="cross-encoder-testing/reranker-bert-tiny-gooaq-bce")
query = "City in Bosnia and Herzegovina"
docs_before_texts = ["Berlin", "Belgrade", "Sarajevo"]
expected_first_text = "Sarajevo"
expected_scores = [0.14568544924259186, 0.18189962208271027, 0.5728498697280884]
docs_before = [Document(content=text) for text in docs_before_texts]
output = ranker.run(query=query, documents=docs_before)
docs_after = output["documents"]
assert len(docs_after) == 3
assert docs_after[0].content == expected_first_text
sorted_scores = sorted(expected_scores, reverse=True)
assert docs_after[0].score == pytest.approx(sorted_scores[0], abs=1e-6)
assert docs_after[1].score == pytest.approx(sorted_scores[1], abs=1e-6)
assert docs_after[2].score == pytest.approx(sorted_scores[2], abs=1e-6)
for doc in docs_after:
assert isinstance(doc.score, float)
@pytest.mark.integration
@pytest.mark.slow
def test_run_top_k(self, del_hf_env_vars):
ranker = SentenceTransformersSimilarityRanker(
model="cross-encoder-testing/reranker-bert-tiny-gooaq-bce", top_k=2
)
query = "City in Bosnia and Herzegovina"
docs_before_texts = ["Berlin", "Belgrade", "Sarajevo"]
expected_first_text = "Sarajevo"
docs_before = [Document(content=text) for text in docs_before_texts]
output = ranker.run(query=query, documents=docs_before)
docs_after = output["documents"]
assert len(docs_after) == 2
assert docs_after[0].content == expected_first_text
sorted_scores = sorted([doc.score for doc in docs_after], reverse=True)
assert [doc.score for doc in docs_after] == sorted_scores
for doc in docs_after:
assert isinstance(doc.score, float)
@pytest.mark.integration
@pytest.mark.slow
def test_run_single_document(self, del_hf_env_vars):
ranker = SentenceTransformersSimilarityRanker(
model="cross-encoder-testing/reranker-bert-tiny-gooaq-bce", device=None
)
docs_before = [Document(content="Berlin")]
output = ranker.run(query="City in Germany", documents=docs_before)
docs_after = output["documents"]
assert len(docs_after) == 1
assert isinstance(docs_after[0].score, float)
test/components/retrievers/test_multi_query_embedding_retriever.py
+18 -25
View File
@@ -10,8 +10,7 @@ import numpy as np
import pytest
from haystack import Document, Pipeline, component
from haystack.components.embedders import SentenceTransformersTextEmbedder
from haystack.components.embedders.sentence_transformers_document_embedder import SentenceTransformersDocumentEmbedder
from haystack.components.embedders import OpenAIDocumentEmbedder, OpenAITextEmbedder
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.query import QueryExpander
from haystack.components.retrievers import InMemoryEmbeddingRetriever, MultiQueryEmbeddingRetriever
@@ -66,7 +65,7 @@ class TestMultiQueryEmbeddingRetriever:
def document_store_with_embeddings(self, sample_documents):
"""Create a document store populated with embedded documents."""
document_store = InMemoryDocumentStore()
doc_embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
doc_embedder = OpenAIDocumentEmbedder()
doc_writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.SKIP)
embedded_docs = doc_embedder.run(sample_documents)["documents"]
@@ -147,7 +146,8 @@ class TestMultiQueryEmbeddingRetriever:
},
}
def test_from_dict(self):
def test_from_dict(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
data = {
"type": "haystack.components.retrievers.multi_query_embedding_retriever.MultiQueryEmbeddingRetriever", # noqa E501
"init_parameters": {
@@ -174,24 +174,18 @@ class TestMultiQueryEmbeddingRetriever:
},
},
"query_embedder": {
"type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", # noqa E501
"type": "haystack.components.embedders.openai_text_embedder.OpenAITextEmbedder",
"init_parameters": {
"model": "sentence-transformers/all-MiniLM-L6-v2",
"token": {"type": "env_var", "env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False},
"api_key": {"env_vars": ["OPENAI_API_KEY"], "strict": True, "type": "env_var"},
"api_base_url": None,
"dimensions": None,
"model": "text-embedding-ada-002",
"organization": None,
"http_client_kwargs": None,
"prefix": "",
"suffix": "",
"batch_size": 32,
"progress_bar": True,
"normalize_embeddings": False,
"trust_remote_code": False,
"local_files_only": False,
"truncate_dim": None,
"model_kwargs": None,
"tokenizer_kwargs": None,
"config_kwargs": None,
"precision": "float32",
"encode_kwargs": None,
"backend": "torch",
"timeout": None,
"max_retries": None,
},
},
"max_workers": 2,
@@ -236,11 +230,11 @@ class TestMultiQueryEmbeddingRetriever:
assert contents.count("Solar energy is renewable") == 1
assert contents.count("Wind energy is clean") == 1
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
@pytest.mark.slow
def test_run_with_filters(self, del_hf_env_vars, document_store_with_embeddings):
def test_run_with_filters(self, document_store_with_embeddings):
in_memory_retriever = InMemoryEmbeddingRetriever(document_store=document_store_with_embeddings)
query_embedder = SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
query_embedder = OpenAITextEmbedder()
multi_retriever = MultiQueryEmbeddingRetriever(retriever=in_memory_retriever, query_embedder=query_embedder)
result = multi_retriever.run(["energy"], {"filters": {"field": "category", "operator": "==", "value": "solar"}})
assert "documents" in result
@@ -251,13 +245,12 @@ class TestMultiQueryEmbeddingRetriever:
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
@pytest.mark.slow
def test_pipeline_integration(self, del_hf_env_vars, document_store_with_embeddings):
def test_pipeline_integration(self, document_store_with_embeddings):
expander = QueryExpander(
chat_generator=OpenAIChatGenerator(model="gpt-4.1-nano"), n_expansions=3, include_original_query=True
)
in_memory_retriever = InMemoryEmbeddingRetriever(document_store=document_store_with_embeddings)
query_embedder = SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
query_embedder = OpenAITextEmbedder()
multiquery_retriever = MultiQueryEmbeddingRetriever(
retriever=in_memory_retriever, query_embedder=query_embedder, max_workers=3
)
test/components/retrievers/test_multi_retriever.py
+16 -11
View File
@@ -2,13 +2,14 @@
#
# SPDX-License-Identifier: Apache-2.0
import os
from typing import Any
from unittest.mock import ANY
import pytest
from haystack import Document, component
from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
from haystack.components.embedders import OpenAIDocumentEmbedder, OpenAITextEmbedder
from haystack.components.retrievers import (
InMemoryBM25Retriever,
InMemoryEmbeddingRetriever,
@@ -75,7 +76,7 @@ def sample_documents():
def document_store_with_embeddings(sample_documents):
"""Create a document store populated with embedded documents."""
document_store = InMemoryDocumentStore()
doc_embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
doc_embedder = OpenAIDocumentEmbedder()
doc_writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.SKIP)
embedded_docs = doc_embedder.run(sample_documents)["documents"]
doc_writer.run(documents=embedded_docs)
@@ -91,7 +92,7 @@ def bm25_retriever(document_store_with_embeddings):
def embedding_retriever(document_store_with_embeddings):
return TextEmbeddingRetriever(
retriever=InMemoryEmbeddingRetriever(document_store=document_store_with_embeddings),
text_embedder=SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"),
text_embedder=OpenAITextEmbedder(),
)
@@ -312,21 +313,24 @@ class TestMultiRetriever:
with pytest.raises(ImportError, match="Could not import class"):
MultiRetriever.from_dict(data)
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
def test_run_with_filters(self, del_hf_env_vars, bm25_retriever, embedding_retriever):
def test_run_with_filters(self, bm25_retriever, embedding_retriever):
retriever = MultiRetriever(retrievers={"bm25": bm25_retriever, "embedding": embedding_retriever})
result = retriever.run(query="energy", filters={"field": "meta.category", "operator": "==", "value": "solar"})
assert len(result["documents"]) == 1
assert result["documents"][0].meta["category"] == "solar"
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
def test_run_with_top_k(self, del_hf_env_vars, bm25_retriever, embedding_retriever):
def test_run_with_top_k(self, bm25_retriever, embedding_retriever):
retriever = MultiRetriever(retrievers={"bm25": bm25_retriever, "embedding": embedding_retriever})
result = retriever.run(query="energy", top_k=2)
assert len(result["documents"]) == 2
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
def test_run_with_active_retrievers_integration(self, del_hf_env_vars, bm25_retriever, embedding_retriever):
def test_run_with_active_retrievers_integration(self, bm25_retriever, embedding_retriever):
retriever = MultiRetriever(retrievers={"bm25": bm25_retriever, "embedding": embedding_retriever})
result_bm25_active = retriever.run(query="energy", active_retrievers=["bm25"])
result_bm25 = bm25_retriever.run(query="energy")
@@ -448,9 +452,10 @@ class TestMultiRetrieverAsync:
assert len(result["documents"]) == 1
assert result["documents"][0].id == "async1"
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
@pytest.mark.asyncio
async def test_run_async_with_filters(self, del_hf_env_vars, bm25_retriever, embedding_retriever):
async def test_run_async_with_filters(self, bm25_retriever, embedding_retriever):
retriever = MultiRetriever(retrievers={"bm25": bm25_retriever, "embedding": embedding_retriever})
result = await retriever.run_async(
query="energy", filters={"field": "meta.category", "operator": "==", "value": "solar"}
@@ -458,18 +463,18 @@ class TestMultiRetrieverAsync:
assert len(result["documents"]) == 1
assert result["documents"][0].meta["category"] == "solar"
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
@pytest.mark.asyncio
async def test_run_async_with_top_k(self, del_hf_env_vars, bm25_retriever, embedding_retriever):
async def test_run_async_with_top_k(self, bm25_retriever, embedding_retriever):
retriever = MultiRetriever(retrievers={"bm25": bm25_retriever, "embedding": embedding_retriever})
result = await retriever.run_async(query="energy", top_k=2)
assert len(result["documents"]) == 2
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
@pytest.mark.asyncio
async def test_run_async_with_active_retrievers_integration(
self, del_hf_env_vars, bm25_retriever, embedding_retriever
):
async def test_run_async_with_active_retrievers_integration(self, bm25_retriever, embedding_retriever):
retriever = MultiRetriever(retrievers={"bm25": bm25_retriever, "embedding": embedding_retriever})
result_bm25_active = await retriever.run_async(query="energy", active_retrievers=["bm25"])
result_bm25 = await bm25_retriever.run_async(query="energy")
test/components/retrievers/test_text_embedding_retriever.py
+21 -25
View File
@@ -2,6 +2,7 @@
#
# SPDX-License-Identifier: Apache-2.0
import os
from typing import Any
from unittest.mock import ANY
@@ -9,7 +10,7 @@ import numpy as np
import pytest
from haystack import Document, component
from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
from haystack.components.embedders import OpenAIDocumentEmbedder, OpenAITextEmbedder
from haystack.components.retrievers import InMemoryEmbeddingRetriever, TextEmbeddingRetriever
from haystack.components.writers import DocumentWriter
from haystack.document_stores.in_memory import InMemoryDocumentStore
@@ -38,7 +39,7 @@ class TestTextEmbeddingRetriever:
def document_store_with_embeddings(self, sample_documents):
"""Create a document store populated with embedded documents."""
document_store = InMemoryDocumentStore()
doc_embedder = SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2")
doc_embedder = OpenAIDocumentEmbedder()
doc_writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.SKIP)
embedded_docs = doc_embedder.run(sample_documents)["documents"]
@@ -118,7 +119,8 @@ class TestTextEmbeddingRetriever:
},
}
def test_from_dict(self):
def test_from_dict(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
data = {
"type": "haystack.components.retrievers.text_embedding_retriever.TextEmbeddingRetriever",
"init_parameters": {
@@ -145,24 +147,18 @@ class TestTextEmbeddingRetriever:
},
},
"text_embedder": {
"type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder", # noqa E501
"type": "haystack.components.embedders.openai_text_embedder.OpenAITextEmbedder",
"init_parameters": {
"model": "sentence-transformers/all-MiniLM-L6-v2",
"token": {"type": "env_var", "env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False},
"api_key": {"env_vars": ["OPENAI_API_KEY"], "strict": True, "type": "env_var"},
"api_base_url": None,
"dimensions": None,
"model": "text-embedding-ada-002",
"organization": None,
"http_client_kwargs": None,
"prefix": "",
"suffix": "",
"batch_size": 32,
"progress_bar": True,
"normalize_embeddings": False,
"trust_remote_code": False,
"local_files_only": False,
"truncate_dim": None,
"model_kwargs": None,
"tokenizer_kwargs": None,
"config_kwargs": None,
"precision": "float32",
"encode_kwargs": None,
"backend": "torch",
"timeout": None,
"max_retries": None,
},
},
},
@@ -170,25 +166,25 @@ class TestTextEmbeddingRetriever:
result = TextEmbeddingRetriever.from_dict(data)
assert isinstance(result, TextEmbeddingRetriever)
assert isinstance(result.retriever, InMemoryEmbeddingRetriever)
assert isinstance(result.text_embedder, SentenceTransformersTextEmbedder)
assert isinstance(result.text_embedder, OpenAITextEmbedder)
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
@pytest.mark.slow
def test_run_with_filters(self, del_hf_env_vars, document_store_with_embeddings):
def test_run_with_filters(self, document_store_with_embeddings):
retriever = TextEmbeddingRetriever(
retriever=InMemoryEmbeddingRetriever(document_store=document_store_with_embeddings),
text_embedder=SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"),
text_embedder=OpenAITextEmbedder(),
)
result = retriever.run(query="energy", filters={"field": "meta.category", "operator": "==", "value": "solar"})
assert "documents" in result
assert all(doc.meta.get("category") == "solar" for doc in result["documents"])
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
@pytest.mark.integration
@pytest.mark.slow
def test_run_with_top_k(self, del_hf_env_vars, document_store_with_embeddings):
def test_run_with_top_k(self, document_store_with_embeddings):
retriever = TextEmbeddingRetriever(
retriever=InMemoryEmbeddingRetriever(document_store=document_store_with_embeddings),
text_embedder=SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"),
text_embedder=OpenAITextEmbedder(),
)
result = retriever.run(query="energy", top_k=2)
assert "documents" in result
test/tools/test_pipeline_tool.py
+32 -27
View File
@@ -7,23 +7,51 @@ from unittest.mock import ANY
import pytest
from haystack import Document, Pipeline
from haystack import Document, Pipeline, component
from haystack.components.agents import Agent
from haystack.components.embedders.openai_document_embedder import OpenAIDocumentEmbedder
from haystack.components.embedders.openai_text_embedder import OpenAITextEmbedder
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.rankers.sentence_transformers_similarity import SentenceTransformersSimilarityRanker
from haystack.components.retrievers import InMemoryBM25Retriever, InMemoryEmbeddingRetriever
from haystack.dataclasses import ChatMessage
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.tools import PipelineTool
@component
class MockSimilarityRanker:
"""Mock ranker used to build a sample pipeline for tests."""
@component.output_types(documents=list[Document])
def run(
self,
documents: list[Document],
query: str,
top_k: int | None = None,
scale_score: bool | None = None,
score_threshold: float | None = None,
) -> dict[str, list[Document]]:
"""
Returns a list of documents ranked by their similarity to the given query.
:param documents: List of documents to rank.
:param query: The input query to compare the documents to.
:param top_k: The maximum number of documents to return.
:param scale_score: If `True`, scales the raw logit predictions using a Sigmoid activation function.
If `False`, disables scaling of the raw logit predictions.
If set, overrides the value set at initialization.
:param score_threshold: Use it to return documents only with a score above this threshold.
If set, overrides the value set at initialization.
"""
ranked = documents[:top_k] if top_k is not None else documents
return {"documents": ranked}
@pytest.fixture
def sample_pipeline():
pipeline = Pipeline()
pipeline.add_component("bm25_retriever", InMemoryBM25Retriever(document_store=InMemoryDocumentStore()))
pipeline.add_component("ranker", SentenceTransformersSimilarityRanker(model="fake-model"))
pipeline.add_component("ranker", MockSimilarityRanker())
pipeline.connect("bm25_retriever", "ranker")
return pipeline
@@ -55,30 +83,7 @@ def sample_pipeline_dict():
"filter_policy": "replace",
},
},
"ranker": {
"type": "haystack.components.rankers.sentence_transformers_similarity."
"SentenceTransformersSimilarityRanker",
"init_parameters": {
"device": {"type": "single", "device": ANY},
"model": "fake-model",
"token": {"type": "env_var", "env_vars": ["HF_API_TOKEN", "HF_TOKEN"], "strict": False},
"top_k": 10,
"query_prefix": "",
"query_suffix": "",
"document_prefix": "",
"document_suffix": "",
"meta_fields_to_embed": [],
"embedding_separator": "\n",
"scale_score": True,
"score_threshold": None,
"trust_remote_code": False,
"model_kwargs": None,
"tokenizer_kwargs": None,
"config_kwargs": None,
"backend": "torch",
"batch_size": 16,
},
},
"ranker": {"type": "test_pipeline_tool.MockSimilarityRanker", "init_parameters": {}},
},
"connections": [{"sender": "bm25_retriever.documents", "receiver": "ranker.documents"}],
"connection_type_validation": True,