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

feat!: Rename model_name or model_name_or_path to model in all Embedder classes (#6733)

Browse Source 
* rename model parameter in the openai doc embedder

* fix tests for openai doc embedder

* rename model parameter in the openai text embedder

* fix tests for openai text embedder

* rename model parameter in the st doc embedder

* fix tests for st doc embedder

* rename model parameter in the st backend

* fix tests for st backend

* rename model parameter in the st text embedder

* fix tests for st text embedder

* fix docstring

* fix pipeline utils

* fix e2e

* reno

* fix the indexing pipeline _create_embedder function

* fix e2e eval rag pipeline

* pytest
This commit is contained in:
ZanSara
2024-01-12 15:30:17 +01:00
committed by GitHub
parent 3156343dce
commit 288ed150c9
22 changed files with 98 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files
e2e/pipelines/test_dense_doc_search.py
+2 -4
View File
@@ -25,8 +25,7 @@ def test_dense_doc_search_pipeline(tmp_path, samples_path):
instance=DocumentSplitter(split_by="sentence", split_length=250, split_overlap=30), name="splitter"
)
indexing_pipeline.add_component(
instance=SentenceTransformersDocumentEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
name="embedder",
instance=SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"), name="embedder"
)
indexing_pipeline.add_component(instance=DocumentWriter(document_store=InMemoryDocumentStore()), name="writer")
@@ -60,8 +59,7 @@ def test_dense_doc_search_pipeline(tmp_path, samples_path):
# Create the querying pipeline
query_pipeline = Pipeline()
query_pipeline.add_component(
instance=SentenceTransformersTextEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
name="text_embedder",
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"
e2e/pipelines/test_eval_dense_doc_search.py
+2 -4
View File
@@ -25,8 +25,7 @@ def test_dense_doc_search_pipeline(samples_path):
instance=DocumentSplitter(split_by="sentence", split_length=250, split_overlap=30), name="splitter"
)
indexing_pipeline.add_component(
instance=SentenceTransformersDocumentEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
name="embedder",
instance=SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"), name="embedder"
)
indexing_pipeline.add_component(instance=DocumentWriter(document_store=InMemoryDocumentStore()), name="writer")
@@ -45,8 +44,7 @@ def test_dense_doc_search_pipeline(samples_path):
# Create the querying pipeline
query_pipeline = Pipeline()
query_pipeline.add_component(
instance=SentenceTransformersTextEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
name="text_embedder",
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"
e2e/pipelines/test_eval_hybrid_doc_search_pipeline.py
+1 -2
View File
@@ -13,8 +13,7 @@ def test_hybrid_doc_search_pipeline():
hybrid_pipeline = Pipeline()
hybrid_pipeline.add_component(instance=InMemoryBM25Retriever(document_store=document_store), name="bm25_retriever")
hybrid_pipeline.add_component(
instance=SentenceTransformersTextEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
name="text_embedder",
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"
e2e/pipelines/test_eval_rag_pipelines.py
+2 -3
View File
@@ -93,8 +93,7 @@ def test_embedding_retrieval_rag_pipeline(tmp_path):
"""
rag_pipeline = Pipeline()
rag_pipeline.add_component(
instance=SentenceTransformersTextEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
name="text_embedder",
instance=SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"), name="text_embedder"
)
rag_pipeline.add_component(
instance=InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore()), name="retriever"
@@ -124,7 +123,7 @@ def test_embedding_retrieval_rag_pipeline(tmp_path):
document_store = rag_pipeline.get_component("retriever").document_store
indexing_pipeline = Pipeline()
indexing_pipeline.add_component(
instance=SentenceTransformersDocumentEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
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")
e2e/pipelines/test_hybrid_doc_search_pipeline.py
+1 -2
View File
@@ -14,8 +14,7 @@ def test_hybrid_doc_search_pipeline(tmp_path):
hybrid_pipeline = Pipeline()
hybrid_pipeline.add_component(instance=InMemoryBM25Retriever(document_store=document_store), name="bm25_retriever")
hybrid_pipeline.add_component(
instance=SentenceTransformersTextEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
name="text_embedder",
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"
e2e/pipelines/test_preprocessing_pipeline.py
+1 -2
View File
@@ -25,8 +25,7 @@ def test_preprocessing_pipeline(tmp_path):
instance=DocumentSplitter(split_by="sentence", split_length=1), name="splitter"
)
preprocessing_pipeline.add_component(
instance=SentenceTransformersDocumentEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
name="embedder",
instance=SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"), 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")
e2e/pipelines/test_rag_pipelines.py
+2 -3
View File
@@ -95,8 +95,7 @@ def test_embedding_retrieval_rag_pipeline(tmp_path):
"""
rag_pipeline = Pipeline()
rag_pipeline.add_component(
instance=SentenceTransformersTextEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
name="text_embedder",
instance=SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"), name="text_embedder"
)
rag_pipeline.add_component(
instance=InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore()), name="retriever"
@@ -131,7 +130,7 @@ def test_embedding_retrieval_rag_pipeline(tmp_path):
document_store = rag_pipeline.get_component("retriever").document_store
indexing_pipeline = Pipeline()
indexing_pipeline.add_component(
instance=SentenceTransformersDocumentEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
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")
examples/pipelines/indexing_pipeline.py
+1 -2
View File
@@ -20,8 +20,7 @@ p.add_component(instance=DocumentJoiner(), name="joiner")
p.add_component(instance=DocumentCleaner(), name="cleaner")
p.add_component(instance=DocumentSplitter(split_by="sentence", split_length=250, split_overlap=30), name="splitter")
p.add_component(
instance=SentenceTransformersDocumentEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
name="embedder",
instance=SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"), name="embedder"
)
p.add_component(instance=DocumentWriter(document_store=InMemoryDocumentStore()), name="writer")
haystack/components/embedders/backends/sentence_transformers_backend.py
+5 -11
View File
@@ -14,15 +14,13 @@ class _SentenceTransformersEmbeddingBackendFactory:
_instances: Dict[str, "_SentenceTransformersEmbeddingBackend"] = {}
@staticmethod
def get_embedding_backend(
model_name_or_path: str, device: Optional[str] = None, use_auth_token: Union[bool, str, None] = None
):
embedding_backend_id = f"{model_name_or_path}{device}{use_auth_token}"
def get_embedding_backend(model: str, device: Optional[str] = None, use_auth_token: Union[bool, str, None] = None):
embedding_backend_id = f"{model}{device}{use_auth_token}"
if embedding_backend_id in _SentenceTransformersEmbeddingBackendFactory._instances:
return _SentenceTransformersEmbeddingBackendFactory._instances[embedding_backend_id]
embedding_backend = _SentenceTransformersEmbeddingBackend(
model_name_or_path=model_name_or_path, device=device, use_auth_token=use_auth_token
model=model, device=device, use_auth_token=use_auth_token
)
_SentenceTransformersEmbeddingBackendFactory._instances[embedding_backend_id] = embedding_backend
return embedding_backend
@@ -33,13 +31,9 @@ class _SentenceTransformersEmbeddingBackend:
Class to manage Sentence Transformers embeddings.
"""
def __init__(
self, model_name_or_path: str, device: Optional[str] = None, use_auth_token: Union[bool, str, None] = None
):
def __init__(self, model: str, device: Optional[str] = None, use_auth_token: Union[bool, str, None] = None):
sentence_transformers_import.check()
self.model = SentenceTransformer(
model_name_or_path=model_name_or_path, device=device, use_auth_token=use_auth_token
)
self.model = SentenceTransformer(model_name_or_path=model, device=device, use_auth_token=use_auth_token)
def embed(self, data: List[str], **kwargs) -> List[List[float]]:
embeddings = self.model.encode(data, **kwargs).tolist()
haystack/components/embedders/openai_document_embedder.py
+6 -6
View File
@@ -31,7 +31,7 @@ class OpenAIDocumentEmbedder:
def __init__(
self,
api_key: Optional[str] = None,
model_name: str = "text-embedding-ada-002",
model: str = "text-embedding-ada-002",
api_base_url: Optional[str] = None,
organization: Optional[str] = None,
prefix: str = "",
@@ -45,7 +45,7 @@ class OpenAIDocumentEmbedder:
Create a OpenAIDocumentEmbedder component.
:param api_key: The OpenAI API key. It can be explicitly provided or automatically read from the
environment variable OPENAI_API_KEY (recommended).
:param model_name: The name of the model to use.
:param model: The name of the model to use.
:param api_base_url: The OpenAI API Base url, defaults to None. For more details, see OpenAI [docs](https://platform.openai.com/docs/api-reference/audio).
:param organization: The Organization ID, defaults to `None`. See
[production best practices](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization).
@@ -57,7 +57,7 @@ class OpenAIDocumentEmbedder:
:param meta_fields_to_embed: List of meta fields that should be embedded along with the Document text.
:param embedding_separator: Separator used to concatenate the meta fields to the Document text.
"""
self.model_name = model_name
self.model = model
self.api_base_url = api_base_url
self.organization = organization
self.prefix = prefix
@@ -73,7 +73,7 @@ class OpenAIDocumentEmbedder:
"""
Data that is sent to Posthog for usage analytics.
"""
return {"model": self.model_name}
return {"model": self.model}
def to_dict(self) -> Dict[str, Any]:
"""
@@ -82,7 +82,7 @@ class OpenAIDocumentEmbedder:
"""
return default_to_dict(
self,
model_name=self.model_name,
model=self.model,
organization=self.organization,
api_base_url=self.api_base_url,
prefix=self.prefix,
@@ -124,7 +124,7 @@ class OpenAIDocumentEmbedder:
range(0, len(texts_to_embed), batch_size), disable=not self.progress_bar, desc="Calculating embeddings"
):
batch = texts_to_embed[i : i + batch_size]
response = self.client.embeddings.create(model=self.model_name, input=batch)
response = self.client.embeddings.create(model=self.model, input=batch)
embeddings = [el.embedding for el in response.data]
all_embeddings.extend(embeddings)
haystack/components/embedders/openai_text_embedder.py
+6 -6
View File
@@ -29,7 +29,7 @@ class OpenAITextEmbedder:
def __init__(
self,
api_key: Optional[str] = None,
model_name: str = "text-embedding-ada-002",
model: str = "text-embedding-ada-002",
api_base_url: Optional[str] = None,
organization: Optional[str] = None,
prefix: str = "",
@@ -40,7 +40,7 @@ class OpenAITextEmbedder:
:param api_key: The OpenAI API key. It can be explicitly provided or automatically read from the
environment variable OPENAI_API_KEY (recommended).
:param model_name: The name of the OpenAI model to use. For more details on the available models,
:param model: The name of the OpenAI model to use. For more details on the available models,
see [OpenAI documentation](https://platform.openai.com/docs/guides/embeddings/embedding-models).
:param organization: The Organization ID, defaults to `None`. See
[production best practices](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization).
@@ -48,7 +48,7 @@ class OpenAITextEmbedder:
:param prefix: A string to add to the beginning of each text.
:param suffix: A string to add to the end of each text.
"""
self.model_name = model_name
self.model = model
self.organization = organization
self.prefix = prefix
self.suffix = suffix
@@ -59,7 +59,7 @@ class OpenAITextEmbedder:
"""
Data that is sent to Posthog for usage analytics.
"""
return {"model": self.model_name}
return {"model": self.model}
def to_dict(self) -> Dict[str, Any]:
"""
@@ -68,7 +68,7 @@ class OpenAITextEmbedder:
"""
return default_to_dict(
self, model_name=self.model_name, organization=self.organization, prefix=self.prefix, suffix=self.suffix
self, model=self.model, organization=self.organization, prefix=self.prefix, suffix=self.suffix
)
@component.output_types(embedding=List[float], meta=Dict[str, Any])
@@ -86,7 +86,7 @@ class OpenAITextEmbedder:
# replace newlines, which can negatively affect performance.
text_to_embed = text_to_embed.replace("\n", " ")
response = self.client.embeddings.create(model=self.model_name, input=text_to_embed)
response = self.client.embeddings.create(model=self.model, input=text_to_embed)
meta = {"model": response.model, "usage": dict(response.usage)}
return {"embedding": response.data[0].embedding, "meta": meta}
haystack/components/embedders/sentence_transformers_document_embedder.py
+6 -6
View File
@@ -29,7 +29,7 @@ class SentenceTransformersDocumentEmbedder:
def __init__(
self,
model_name_or_path: str = "sentence-transformers/all-mpnet-base-v2",
model: str = "sentence-transformers/all-mpnet-base-v2",
device: Optional[str] = None,
token: Union[bool, str, None] = None,
prefix: str = "",
@@ -43,7 +43,7 @@ class SentenceTransformersDocumentEmbedder:
"""
Create a SentenceTransformersDocumentEmbedder component.
:param model_name_or_path: Local path or name of the model in Hugging Face's model hub,
:param model: Local path or name of the model in Hugging Face's model hub,
such as ``'sentence-transformers/all-mpnet-base-v2'``.
:param device: Device (like 'cuda' / 'cpu') that should be used for computation.
Defaults to CPU.
@@ -61,7 +61,7 @@ class SentenceTransformersDocumentEmbedder:
:param embedding_separator: Separator used to concatenate the meta fields to the Document content.
"""
self.model_name_or_path = model_name_or_path
self.model = model
# TODO: remove device parameter and use Haystack's device management once migrated
self.device = device or "cpu"
self.token = token
@@ -77,7 +77,7 @@ class SentenceTransformersDocumentEmbedder:
"""
Data that is sent to Posthog for usage analytics.
"""
return {"model": self.model_name_or_path}
return {"model": self.model}
def to_dict(self) -> Dict[str, Any]:
"""
@@ -85,7 +85,7 @@ class SentenceTransformersDocumentEmbedder:
"""
return default_to_dict(
self,
model_name_or_path=self.model_name_or_path,
model=self.model,
device=self.device,
token=self.token if not isinstance(self.token, str) else None, # don't serialize valid tokens
prefix=self.prefix,
@@ -103,7 +103,7 @@ class SentenceTransformersDocumentEmbedder:
"""
if not hasattr(self, "embedding_backend"):
self.embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model_name_or_path=self.model_name_or_path, device=self.device, use_auth_token=self.token
model=self.model, device=self.device, use_auth_token=self.token
)
@component.output_types(documents=List[Document])
haystack/components/embedders/sentence_transformers_text_embedder.py
+6 -6
View File
@@ -28,7 +28,7 @@ class SentenceTransformersTextEmbedder:
def __init__(
self,
model_name_or_path: str = "sentence-transformers/all-mpnet-base-v2",
model: str = "sentence-transformers/all-mpnet-base-v2",
device: Optional[str] = None,
token: Union[bool, str, None] = None,
prefix: str = "",
@@ -40,7 +40,7 @@ class SentenceTransformersTextEmbedder:
"""
Create a SentenceTransformersTextEmbedder component.
:param model_name_or_path: Local path or name of the model in Hugging Face's model hub,
:param model: Local path or name of the model in Hugging Face's model hub,
such as ``'sentence-transformers/all-mpnet-base-v2'``.
:param device: Device (like 'cuda' / 'cpu') that should be used for computation.
Defaults to CPU.
@@ -56,7 +56,7 @@ class SentenceTransformersTextEmbedder:
:param normalize_embeddings: If set to true, returned vectors will have length 1.
"""
self.model_name_or_path = model_name_or_path
self.model = model
# TODO: remove device parameter and use Haystack's device management once migrated
self.device = device or "cpu"
self.token = token
@@ -70,7 +70,7 @@ class SentenceTransformersTextEmbedder:
"""
Data that is sent to Posthog for usage analytics.
"""
return {"model": self.model_name_or_path}
return {"model": self.model}
def to_dict(self) -> Dict[str, Any]:
"""
@@ -78,7 +78,7 @@ class SentenceTransformersTextEmbedder:
"""
return default_to_dict(
self,
model_name_or_path=self.model_name_or_path,
model=self.model,
device=self.device,
token=self.token if not isinstance(self.token, str) else None, # don't serialize valid tokens
prefix=self.prefix,
@@ -94,7 +94,7 @@ class SentenceTransformersTextEmbedder:
"""
if not hasattr(self, "embedding_backend"):
self.embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model_name_or_path=self.model_name_or_path, device=self.device, use_auth_token=self.token
model=self.model, device=self.device, use_auth_token=self.token
)
@component.output_types(embedding=List[float])
haystack/components/joiners/document_joiner.py
+1 -1
View File
@@ -31,7 +31,7 @@ class DocumentJoiner:
p = Pipeline()
p.add_component(instance=InMemoryBM25Retriever(document_store=document_store), name="bm25_retriever")
p.add_component(
instance=SentenceTransformersTextEmbedder(model_name_or_path="sentence-transformers/all-MiniLM-L6-v2"),
instance=SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"),
name="text_embedder",
)
p.add_component(instance=InMemoryEmbeddingRetriever(document_store=document_store), name="embedding_retriever")
haystack/pipeline_utils/indexing.py
+4 -17
View File
@@ -1,4 +1,3 @@
import inspect
import os
import re
from pathlib import Path
@@ -197,22 +196,10 @@ class _IndexingPipeline:
)
return self._create_embedder(embedder_class, embedding_model, init_kwargs)
def _create_embedder(
self, embedder_class: Type, model_name: str, init_kwargs: Optional[Dict[str, Any]] = None
) -> Any:
init_signature = inspect.signature(embedder_class.__init__)
kwargs = {**(init_kwargs or {})}
# Determine the correct parameter name and set it
if "model_name_or_path" in init_signature.parameters:
kwargs["model_name_or_path"] = model_name
elif "model_name" in init_signature.parameters:
kwargs["model_name"] = model_name
else:
raise ValueError(f"Could not find a parameter for the model name in the embedder class {embedder_class}")
# Instantiate the class
def _create_embedder(self, embedder_class: Type, model: str, init_kwargs: Optional[Dict[str, Any]] = None) -> Any:
# Note: here we assume the embedder accepts a parameter called `model` that takes the model's name or path.
# See https://github.com/deepset-ai/haystack/issues/6534
kwargs = {**(init_kwargs or {}), "model": model}
return embedder_class(**kwargs)
def _list_files_recursively(self, path: Union[str, Path]) -> List[str]:
haystack/pipeline_utils/rag.py
+1 -1
View File
@@ -92,7 +92,7 @@ def resolve_embedder(embedding_model: str) -> SentenceTransformersTextEmbedder:
:param embedding_model: The embedding model to use.
"""
try:
embedder = SentenceTransformersTextEmbedder(model_name_or_path=embedding_model)
embedder = SentenceTransformersTextEmbedder(model=embedding_model)
except Exception:
raise ValueError(
f"Embedding model: {embedding_model} is not supported. Please provide a SentenceTransformers model."
releasenotes/notes/rename-model-param-embedders-7cc87a768554724d.yaml
+3
View File
@@ -0,0 +1,3 @@
---
upgrade:
- Rename the embedder parameters `model_name` and `model_name_or_path` to `model`. This change affects all Embedder classes.
test/components/embedders/test_openai_document_embedder.py
+7 -7
View File
@@ -26,7 +26,7 @@ class TestOpenAIDocumentEmbedder:
def test_init_default(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
embedder = OpenAIDocumentEmbedder()
assert embedder.model_name == "text-embedding-ada-002"
assert embedder.model == "text-embedding-ada-002"
assert embedder.organization is None
assert embedder.prefix == ""
assert embedder.suffix == ""
@@ -38,7 +38,7 @@ class TestOpenAIDocumentEmbedder:
def test_init_with_parameters(self):
embedder = OpenAIDocumentEmbedder(
api_key="fake-api-key",
model_name="model",
model="model",
organization="my-org",
prefix="prefix",
suffix="suffix",
@@ -48,7 +48,7 @@ class TestOpenAIDocumentEmbedder:
embedding_separator=" | ",
)
assert embedder.organization == "my-org"
assert embedder.model_name == "model"
assert embedder.model == "model"
assert embedder.prefix == "prefix"
assert embedder.suffix == "suffix"
assert embedder.batch_size == 64
@@ -68,7 +68,7 @@ class TestOpenAIDocumentEmbedder:
"type": "haystack.components.embedders.openai_document_embedder.OpenAIDocumentEmbedder",
"init_parameters": {
"api_base_url": None,
"model_name": "text-embedding-ada-002",
"model": "text-embedding-ada-002",
"organization": None,
"prefix": "",
"suffix": "",
@@ -82,7 +82,7 @@ class TestOpenAIDocumentEmbedder:
def test_to_dict_with_custom_init_parameters(self):
component = OpenAIDocumentEmbedder(
api_key="fake-api-key",
model_name="model",
model="model",
organization="my-org",
prefix="prefix",
suffix="suffix",
@@ -96,7 +96,7 @@ class TestOpenAIDocumentEmbedder:
"type": "haystack.components.embedders.openai_document_embedder.OpenAIDocumentEmbedder",
"init_parameters": {
"api_base_url": None,
"model_name": "model",
"model": "model",
"organization": "my-org",
"prefix": "prefix",
"suffix": "suffix",
@@ -174,7 +174,7 @@ class TestOpenAIDocumentEmbedder:
model = "text-embedding-ada-002"
embedder = OpenAIDocumentEmbedder(model_name=model, meta_fields_to_embed=["topic"], embedding_separator=" | ")
embedder = OpenAIDocumentEmbedder(model=model, meta_fields_to_embed=["topic"], embedding_separator=" | ")
result = embedder.run(documents=docs)
documents_with_embeddings = result["documents"]
test/components/embedders/test_openai_text_embedder.py
+7 -20
View File
@@ -12,21 +12,17 @@ class TestOpenAITextEmbedder:
embedder = OpenAITextEmbedder()
assert embedder.client.api_key == "fake-api-key"
assert embedder.model_name == "text-embedding-ada-002"
assert embedder.model == "text-embedding-ada-002"
assert embedder.organization is None
assert embedder.prefix == ""
assert embedder.suffix == ""
def test_init_with_parameters(self):
embedder = OpenAITextEmbedder(
api_key="fake-api-key",
model_name="model",
organization="fake-organization",
prefix="prefix",
suffix="suffix",
api_key="fake-api-key", model="model", organization="fake-organization", prefix="prefix", suffix="suffix"
)
assert embedder.client.api_key == "fake-api-key"
assert embedder.model_name == "model"
assert embedder.model == "model"
assert embedder.organization == "fake-organization"
assert embedder.prefix == "prefix"
assert embedder.suffix == "suffix"
@@ -41,27 +37,18 @@ class TestOpenAITextEmbedder:
data = component.to_dict()
assert data == {
"type": "haystack.components.embedders.openai_text_embedder.OpenAITextEmbedder",
"init_parameters": {
"model_name": "text-embedding-ada-002",
"organization": None,
"prefix": "",
"suffix": "",
},
"init_parameters": {"model": "text-embedding-ada-002", "organization": None, "prefix": "", "suffix": ""},
}
def test_to_dict_with_custom_init_parameters(self):
component = OpenAITextEmbedder(
api_key="fake-api-key",
model_name="model",
organization="fake-organization",
prefix="prefix",
suffix="suffix",
api_key="fake-api-key", model="model", organization="fake-organization", prefix="prefix", suffix="suffix"
)
data = component.to_dict()
assert data == {
"type": "haystack.components.embedders.openai_text_embedder.OpenAITextEmbedder",
"init_parameters": {
"model_name": "model",
"model": "model",
"organization": "fake-organization",
"prefix": "prefix",
"suffix": "suffix",
@@ -81,7 +68,7 @@ class TestOpenAITextEmbedder:
def test_run(self):
model = "text-embedding-ada-002"
embedder = OpenAITextEmbedder(model_name=model, prefix="prefix ", suffix=" suffix")
embedder = OpenAITextEmbedder(model=model, prefix="prefix ", suffix=" suffix")
result = embedder.run(text="The food was delicious")
assert len(result["embedding"]) == 1536
test/components/embedders/test_sentence_transformers_document_embedder.py
+15 -17
View File
@@ -8,8 +8,8 @@ from haystack.components.embedders.sentence_transformers_document_embedder impor
class TestSentenceTransformersDocumentEmbedder:
def test_init_default(self):
embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model")
assert embedder.model_name_or_path == "model"
embedder = SentenceTransformersDocumentEmbedder(model="model")
assert embedder.model == "model"
assert embedder.device == "cpu"
assert embedder.token is None
assert embedder.prefix == ""
@@ -22,7 +22,7 @@ class TestSentenceTransformersDocumentEmbedder:
def test_init_with_parameters(self):
embedder = SentenceTransformersDocumentEmbedder(
model_name_or_path="model",
model="model",
device="cuda",
token=True,
prefix="prefix",
@@ -33,7 +33,7 @@ class TestSentenceTransformersDocumentEmbedder:
meta_fields_to_embed=["test_field"],
embedding_separator=" | ",
)
assert embedder.model_name_or_path == "model"
assert embedder.model == "model"
assert embedder.device == "cuda"
assert embedder.token is True
assert embedder.prefix == "prefix"
@@ -45,12 +45,12 @@ class TestSentenceTransformersDocumentEmbedder:
assert embedder.embedding_separator == " | "
def test_to_dict(self):
component = SentenceTransformersDocumentEmbedder(model_name_or_path="model")
component = SentenceTransformersDocumentEmbedder(model="model")
data = component.to_dict()
assert data == {
"type": "haystack.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder",
"init_parameters": {
"model_name_or_path": "model",
"model": "model",
"device": "cpu",
"token": None,
"prefix": "",
@@ -65,7 +65,7 @@ class TestSentenceTransformersDocumentEmbedder:
def test_to_dict_with_custom_init_parameters(self):
component = SentenceTransformersDocumentEmbedder(
model_name_or_path="model",
model="model",
device="cuda",
token="the-token",
prefix="prefix",
@@ -81,7 +81,7 @@ class TestSentenceTransformersDocumentEmbedder:
assert data == {
"type": "haystack.components.embedders.sentence_transformers_document_embedder.SentenceTransformersDocumentEmbedder",
"init_parameters": {
"model_name_or_path": "model",
"model": "model",
"device": "cuda",
"token": None, # the token is not serialized
"prefix": "prefix",
@@ -98,25 +98,23 @@ class TestSentenceTransformersDocumentEmbedder:
"haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory"
)
def test_warmup(self, mocked_factory):
embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model")
embedder = SentenceTransformersDocumentEmbedder(model="model")
mocked_factory.get_embedding_backend.assert_not_called()
embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model_name_or_path="model", device="cpu", use_auth_token=None
)
mocked_factory.get_embedding_backend.assert_called_once_with(model="model", device="cpu", use_auth_token=None)
@patch(
"haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory"
)
def test_warmup_doesnt_reload(self, mocked_factory):
embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model")
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_name_or_path="model")
embedder = SentenceTransformersDocumentEmbedder(model="model")
embedder.embedding_backend = MagicMock()
embedder.embedding_backend.embed = lambda x, **kwargs: np.random.rand(len(x), 16).tolist()
@@ -132,7 +130,7 @@ class TestSentenceTransformersDocumentEmbedder:
assert isinstance(doc.embedding[0], float)
def test_run_wrong_input_format(self):
embedder = SentenceTransformersDocumentEmbedder(model_name_or_path="model")
embedder = SentenceTransformersDocumentEmbedder(model="model")
string_input = "text"
list_integers_input = [1, 2, 3]
@@ -149,7 +147,7 @@ class TestSentenceTransformersDocumentEmbedder:
def test_embed_metadata(self):
embedder = SentenceTransformersDocumentEmbedder(
model_name_or_path="model", meta_fields_to_embed=["meta_field"], embedding_separator="\n"
model="model", meta_fields_to_embed=["meta_field"], embedding_separator="\n"
)
embedder.embedding_backend = MagicMock()
@@ -172,7 +170,7 @@ class TestSentenceTransformersDocumentEmbedder:
def test_prefix_suffix(self):
embedder = SentenceTransformersDocumentEmbedder(
model_name_or_path="model",
model="model",
prefix="my_prefix ",
suffix=" my_suffix",
meta_fields_to_embed=["meta_field"],
test/components/embedders/test_sentence_transformers_embedding_backend.py
+4 -4
View File
@@ -8,11 +8,11 @@ from haystack.components.embedders.backends.sentence_transformers_backend import
@patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer")
def test_factory_behavior(mock_sentence_transformer):
embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model_name_or_path="my_model", device="cpu"
model="my_model", device="cpu"
)
same_embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend("my_model", "cpu")
another_embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model_name_or_path="another_model", device="cpu"
model="another_model", device="cpu"
)
assert same_embedding_backend is embedding_backend
@@ -22,7 +22,7 @@ def test_factory_behavior(mock_sentence_transformer):
@patch("haystack.components.embedders.backends.sentence_transformers_backend.SentenceTransformer")
def test_model_initialization(mock_sentence_transformer):
_SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
model_name_or_path="model", device="cpu", use_auth_token="my_token"
model="model", device="cpu", use_auth_token="my_token"
)
mock_sentence_transformer.assert_called_once_with(
model_name_or_path="model", device="cpu", use_auth_token="my_token"
@@ -31,7 +31,7 @@ def test_model_initialization(mock_sentence_transformer):
@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_name_or_path="model")
embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(model="model")
data = ["sentence1", "sentence2"]
embedding_backend.embed(data=data, normalize_embeddings=True)
test/components/embedders/test_sentence_transformers_text_embedder.py
+15 -17
View File
@@ -8,8 +8,8 @@ from haystack.components.embedders.sentence_transformers_text_embedder import Se
class TestSentenceTransformersTextEmbedder:
def test_init_default(self):
embedder = SentenceTransformersTextEmbedder(model_name_or_path="model")
assert embedder.model_name_or_path == "model"
embedder = SentenceTransformersTextEmbedder(model="model")
assert embedder.model == "model"
assert embedder.device == "cpu"
assert embedder.token is None
assert embedder.prefix == ""
@@ -20,7 +20,7 @@ class TestSentenceTransformersTextEmbedder:
def test_init_with_parameters(self):
embedder = SentenceTransformersTextEmbedder(
model_name_or_path="model",
model="model",
device="cuda",
token=True,
prefix="prefix",
@@ -29,7 +29,7 @@ class TestSentenceTransformersTextEmbedder:
progress_bar=False,
normalize_embeddings=True,
)
assert embedder.model_name_or_path == "model"
assert embedder.model == "model"
assert embedder.device == "cuda"
assert embedder.token is True
assert embedder.prefix == "prefix"
@@ -39,12 +39,12 @@ class TestSentenceTransformersTextEmbedder:
assert embedder.normalize_embeddings is True
def test_to_dict(self):
component = SentenceTransformersTextEmbedder(model_name_or_path="model")
component = SentenceTransformersTextEmbedder(model="model")
data = component.to_dict()
assert data == {
"type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder",
"init_parameters": {
"model_name_or_path": "model",
"model": "model",
"device": "cpu",
"token": None,
"prefix": "",
@@ -57,7 +57,7 @@ class TestSentenceTransformersTextEmbedder:
def test_to_dict_with_custom_init_parameters(self):
component = SentenceTransformersTextEmbedder(
model_name_or_path="model",
model="model",
device="cuda",
token=True,
prefix="prefix",
@@ -70,7 +70,7 @@ class TestSentenceTransformersTextEmbedder:
assert data == {
"type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder",
"init_parameters": {
"model_name_or_path": "model",
"model": "model",
"device": "cuda",
"token": True,
"prefix": "prefix",
@@ -82,12 +82,12 @@ class TestSentenceTransformersTextEmbedder:
}
def test_to_dict_not_serialize_token(self):
component = SentenceTransformersTextEmbedder(model_name_or_path="model", token="awesome-token")
component = SentenceTransformersTextEmbedder(model="model", token="awesome-token")
data = component.to_dict()
assert data == {
"type": "haystack.components.embedders.sentence_transformers_text_embedder.SentenceTransformersTextEmbedder",
"init_parameters": {
"model_name_or_path": "model",
"model": "model",
"device": "cpu",
"token": None,
"prefix": "",
@@ -102,25 +102,23 @@ class TestSentenceTransformersTextEmbedder:
"haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
)
def test_warmup(self, mocked_factory):
embedder = SentenceTransformersTextEmbedder(model_name_or_path="model")
embedder = SentenceTransformersTextEmbedder(model="model")
mocked_factory.get_embedding_backend.assert_not_called()
embedder.warm_up()
mocked_factory.get_embedding_backend.assert_called_once_with(
model_name_or_path="model", device="cpu", use_auth_token=None
)
mocked_factory.get_embedding_backend.assert_called_once_with(model="model", device="cpu", use_auth_token=None)
@patch(
"haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
)
def test_warmup_doesnt_reload(self, mocked_factory):
embedder = SentenceTransformersTextEmbedder(model_name_or_path="model")
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_name_or_path="model")
embedder = SentenceTransformersTextEmbedder(model="model")
embedder.embedding_backend = MagicMock()
embedder.embedding_backend.embed = lambda x, **kwargs: np.random.rand(len(x), 16).tolist()
@@ -133,7 +131,7 @@ class TestSentenceTransformersTextEmbedder:
assert all(isinstance(el, float) for el in embedding)
def test_run_wrong_input_format(self):
embedder = SentenceTransformersTextEmbedder(model_name_or_path="model")
embedder = SentenceTransformersTextEmbedder(model="model")
embedder.embedding_backend = MagicMock()
list_integers_input = [1, 2, 3]