From c1ebdfb734c5f43b193927e2e6600a2a40b2a8cb Mon Sep 17 00:00:00 2001 From: NirantK Date: Tue, 11 Jul 2023 19:33:06 +0530 Subject: [PATCH] * refactor(qdrant_mixin.py): * - Move embedding model initialization to a separate parameter in the methods `index_documents` and `search_documents` * - Remove unused import `OpenAIEmbedding` * - Remove unused import `models` from `qdrant_client.http` * - Remove unused import --- fastvector/qdrant_mixin.py | 25 +++++++++---------------- 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/fastvector/qdrant_mixin.py b/fastvector/qdrant_mixin.py index f258b44..f4cfedd 100644 --- a/fastvector/qdrant_mixin.py +++ b/fastvector/qdrant_mixin.py @@ -1,6 +1,7 @@ from warnings import warn from qdrant_client.models import PointStruct from qdrant_client.http import models +from .embedding import SentenceTransformersEmbedding, OpenAIEmbedding class QdrantClientMixin: @@ -10,15 +11,10 @@ class QdrantClientMixin: docs, batch_size=512, wait=True, - model_name="sentence-transformers/all-MiniLM-L6-v2", + embedding_model=None, ): - # Initialize the SentenceTransformer model - try: - from sentence_transformers import SentenceTransformer - except ImportError: - raise ImportError("Please install the sentence-transformers package to use this method.") - model = SentenceTransformer(model_name) - + if embedding_model is None: + embedding_model = SentenceTransformersEmbedding() n = len(docs["documents"]) if "ids" not in docs: @@ -31,7 +27,7 @@ class QdrantClientMixin: batch_ids = docs["ids"][i : i + batch_size] # Tokenize, embed, and index each document - embeddings = model.encode(batch_docs) + embeddings = embedding_model.encode(batch_docs) # Create a PointStruct for each document points = [ @@ -57,17 +53,14 @@ class QdrantClientMixin: n_results=2, query_filter=None, search_params=None, + embedding_model=None, **kwargs, ): - try: - from sentence_transformers import SentenceTransformer - except ImportError: - raise ImportError("Please install the sentence-transformers package to use this method.") - # Initialize the SentenceTransformer model - model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2") + if embedding_model is None: + embedding_model = SentenceTransformersEmbedding() # Encode the query text - query_vector = model.encode(query_texts)[0] + query_vector = embedding_model.encode(query_texts)[0] # Define default search parameters if not provided if search_params is None: