bug: fix embedding_dim mismatch in DocumentStore (#3183)

* match index dim with embed dim (#3090)

* aligned messages across all docstores

* aligned messages across all docstores (#3090)

* aligned messages across all docstores (#3090)

Author: kalki7

haystack/document_stores/elasticsearch.py

@@ -1479,14 +1479,16 @@ def update_embeddings(
             for result_batch in get_batches_from_generator(result, batch_size):
                 document_batch = [self._convert_es_hit_to_document(hit, return_embedding=False) for hit in result_batch]
                 embeddings = retriever.embed_documents(document_batch)  # type: ignore
-                assert len(document_batch) == len(embeddings)
-
+                if len(document_batch) != len(embeddings):
+                    raise DocumentStoreError(
+                        "The number of embeddings does not match the number of documents in the batch "
+                        f"({len(embeddings)} != {len(document_batch)})"
+                    )
                 if embeddings[0].shape[0] != self.embedding_dim:
                     raise RuntimeError(
-                        f"Embedding dim. of model ({embeddings[0].shape[0]})"
-                        f" doesn't match embedding dim. in DocumentStore ({self.embedding_dim})."
-                        "Specify the arg `embedding_dim` when initializing ElasticsearchDocumentStore()"
+                        f"Embedding dimensions of the model ({embeddings[0].shape[0]}) doesn't match the embedding dimensions of the document store ({self.embedding_dim}). Please reinitiate ElasticsearchDocumentStore() with arg embedding_dim={embeddings[0].shape[0]}."
                     )
+
                 doc_updates = []
                 for doc, emb in zip(document_batch, embeddings):
                     update = {

haystack/document_stores/faiss.py

@@ -22,6 +22,7 @@
 
 from haystack.schema import Document
 from haystack.document_stores.base import get_batches_from_generator
+from haystack.errors import DocumentStoreError
 
 if TYPE_CHECKING:
     from haystack.nodes.retriever import BaseRetriever
@@ -361,7 +362,15 @@ def update_embeddings(
         ) as progress_bar:
             for document_batch in batched_documents:
                 embeddings = retriever.embed_documents(document_batch)  # type: ignore
-                assert len(document_batch) == len(embeddings)
+                if len(document_batch) != len(embeddings):
+                    raise DocumentStoreError(
+                        "The number of embeddings does not match the number of documents in the batch "
+                        f"({len(embeddings)} != {len(document_batch)})"
+                    )
+                if embeddings[0].shape[0] != self.embedding_dim:
+                    raise RuntimeError(
+                        f"Embedding dimensions of the model ({embeddings[0].shape[0]}) doesn't match the embedding dimensions of the document store ({self.embedding_dim}). Please reinitiate FAISSDocumentStore() with arg embedding_dim={embeddings[0].shape[0]}."
+                    )
 
                 embeddings_to_index = np.array(embeddings, dtype="float32")
 

haystack/document_stores/memory.py

@@ -459,16 +459,14 @@ def update_embeddings(
         ) as progress_bar:
             for document_batch in batched_documents:
                 embeddings = retriever.embed_documents(document_batch)  # type: ignore
-                if not len(document_batch) == len(embeddings):
+                if len(document_batch) != len(embeddings):
                     raise DocumentStoreError(
                         "The number of embeddings does not match the number of documents in the batch "
                         f"({len(embeddings)} != {len(document_batch)})"
                     )
                 if embeddings[0].shape[0] != self.embedding_dim:
                     raise RuntimeError(
-                        f"Embedding dim. of model ({embeddings[0].shape[0]})"
-                        f" doesn't match embedding dim. in DocumentStore ({self.embedding_dim})."
-                        "Specify the arg `embedding_dim` when initializing InMemoryDocumentStore()"
+                        f"Embedding dimensions of the model ({embeddings[0].shape[0]}) doesn't match the embedding dimensions of the document store ({self.embedding_dim}). Please reinitiate InMemoryDocumentStore() with arg embedding_dim={embeddings[0].shape[0]}."
                     )
 
                 for doc, emb in zip(document_batch, embeddings):

haystack/document_stores/milvus1.py

@@ -15,6 +15,7 @@
 
 from haystack.schema import Document
 from haystack.document_stores.base import get_batches_from_generator
+from haystack.errors import DocumentStoreError
 
 if TYPE_CHECKING:
     from haystack.nodes.retriever import BaseRetriever
@@ -334,6 +335,16 @@ def update_embeddings(
                 self._delete_vector_ids_from_milvus(documents=document_batch, index=index)
 
                 embeddings = retriever.embed_documents(document_batch)  # type: ignore
+                if len(document_batch) != len(embeddings):
+                    raise DocumentStoreError(
+                        "The number of embeddings does not match the number of documents in the batch "
+                        f"({len(embeddings)} != {len(document_batch)})"
+                    )
+                if embeddings[0].shape[0] != self.embedding_dim:
+                    raise RuntimeError(
+                        f"Embedding dimensions of the model ({embeddings[0].shape[0]}) doesn't match the embedding dimensions of the document store ({self.embedding_dim}). Please reinitiate MilvusDocumentStore() with arg embedding_dim={embeddings[0].shape[0]}."
+                    )
+
                 if self.similarity == "cosine":
                     for embedding in embeddings:
                         self.normalize_embedding(embedding)

haystack/document_stores/milvus2.py

@@ -18,6 +18,7 @@
 from haystack.schema import Document
 from haystack.document_stores.sql import SQLDocumentStore
 from haystack.document_stores.base import get_batches_from_generator
+from haystack.errors import DocumentStoreError
 
 if TYPE_CHECKING:
     from haystack.nodes.retriever.base import BaseRetriever
@@ -369,6 +370,16 @@ def update_embeddings(
                 self._delete_vector_ids_from_milvus(documents=document_batch, index=index)
 
                 embeddings = retriever.embed_documents(document_batch)  # type: ignore
+                if len(document_batch) != len(embeddings):
+                    raise DocumentStoreError(
+                        "The number of embeddings does not match the number of documents in the batch "
+                        f"({len(embeddings)} != {len(document_batch)})"
+                    )
+                if embeddings[0].shape[0] != self.embedding_dim:
+                    raise RuntimeError(
+                        f"Embedding dimensions of the model ({embeddings[0].shape[0]}) doesn't match the embedding dimensions of the document store ({self.embedding_dim}). Please reinitiate MilvusDocumentStore() with arg embedding_dim={embeddings[0].shape[0]}."
+                    )
+
                 if self.cosine:
                     embeddings = [embedding / np.linalg.norm(embedding) for embedding in embeddings]
                 embeddings_list = [embedding.tolist() for embedding in embeddings]

haystack/document_stores/pinecone.py

@@ -11,7 +11,7 @@
 from haystack.document_stores import BaseDocumentStore
 
 from haystack.document_stores.filter_utils import LogicalFilterClause
-from haystack.errors import PineconeDocumentStoreError, DuplicateDocumentError
+from haystack.errors import PineconeDocumentStoreError, DuplicateDocumentError, DocumentStoreError
 
 if TYPE_CHECKING:
     from haystack.nodes.retriever import BaseRetriever
@@ -490,7 +490,15 @@ def update_embeddings(
             for _ in range(0, document_count, batch_size):
                 document_batch = list(islice(documents, batch_size))
                 embeddings = retriever.embed_documents(document_batch)  # type: ignore
-                assert len(document_batch) == len(embeddings)
+                if len(document_batch) != len(embeddings):
+                    raise DocumentStoreError(
+                        "The number of embeddings does not match the number of documents in the batch "
+                        f"({len(embeddings)} != {len(document_batch)})"
+                    )
+                if embeddings[0].shape[0] != self.embedding_dim:
+                    raise RuntimeError(
+                        f"Embedding dimensions of the model ({embeddings[0].shape[0]}) doesn't match the embedding dimensions of the document store ({self.embedding_dim}). Please reinitiate PineconeDocumentStore() with arg embedding_dim={embeddings[0].shape[0]}."
+                    )
 
                 embeddings_to_index = np.array(embeddings, dtype="float32")
                 if self.similarity == "cosine":

haystack/document_stores/weaviate.py

@@ -1228,14 +1228,16 @@ def update_embeddings(
                 self._convert_weaviate_result_to_document(hit, return_embedding=False) for hit in result_batch
             ]
             embeddings = retriever.embed_documents(document_batch)  # type: ignore
-            assert len(document_batch) == len(embeddings)
-
+            if len(document_batch) != len(embeddings):
+                raise DocumentStoreError(
+                    "The number of embeddings does not match the number of documents in the batch "
+                    f"({len(embeddings)} != {len(document_batch)})"
+                )
             if embeddings[0].shape[0] != self.embedding_dim:
                 raise RuntimeError(
-                    f"Embedding dim. of model ({embeddings[0].shape[0]})"
-                    f" doesn't match embedding dim. in DocumentStore ({self.embedding_dim})."
-                    "Specify the arg `embedding_dim` when initializing WeaviateDocumentStore()"
+                    f"Embedding dimensions of the model ({embeddings[0].shape[0]}) doesn't match the embedding dimensions of the document store ({self.embedding_dim}). Please reinitiate WeaviateDocumentStore() with arg embedding_dim={embeddings[0].shape[0]}."
                 )
+
             for doc, emb in zip(document_batch, embeddings):
                 # Using update method to only update the embeddings, other properties will be in tact
                 if self.similarity == "cosine":