diff --git a/qdrant_client/async_qdrant_fastembed.py b/qdrant_client/async_qdrant_fastembed.py index 7faf9108..0a0d8e1b 100644 --- a/qdrant_client/async_qdrant_fastembed.py +++ b/qdrant_client/async_qdrant_fastembed.py @@ -9,41 +9,27 @@ # # ****** WARNING: THIS FILE IS AUTOGENERATED ****** -import uuid -from itertools import tee from typing import Any, Iterable, Sequence, get_args from copy import deepcopy import numpy as np from pydantic import BaseModel from qdrant_client import grpc -from qdrant_client.common.client_warnings import show_warning, show_warning_once from qdrant_client.async_client_base import AsyncQdrantBase -from qdrant_client.embed.embedder import Embedder from qdrant_client.embed.model_embedder import ModelEmbedder from qdrant_client.http import models from qdrant_client.conversions import common_types as types from qdrant_client.conversions.conversion import GrpcToRest from qdrant_client.embed.common import INFERENCE_OBJECT_TYPES from qdrant_client.embed.schema_parser import ModelSchemaParser -from qdrant_client.hybrid.fusion import reciprocal_rank_fusion -from qdrant_client.fastembed_common import FastEmbedMisc, OnnxProvider -from qdrant_client.fastembed_common import ( - QueryResponse, - TextEmbedding, - SparseTextEmbedding, - IDF_EMBEDDING_MODELS, -) +from qdrant_client.fastembed_common import FastEmbedMisc class AsyncQdrantFastembedMixin(AsyncQdrantBase): - DEFAULT_EMBEDDING_MODEL = "BAAI/bge-small-en" DEFAULT_BATCH_SIZE = 8 _FASTEMBED_INSTALLED: bool def __init__(self, parser: ModelSchemaParser, is_local_mode: bool): self.__class__._FASTEMBED_INSTALLED = FastEmbedMisc.is_installed() - self._embedding_model_name: str | None = None - self._sparse_embedding_model_name: str | None = None self._model_embedder = ModelEmbedder(parser=parser, is_local_mode=is_local_mode) super().__init__() @@ -94,124 +80,6 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase): """ return FastEmbedMisc.list_sparse_models() - @property - def embedding_model_name(self) -> str: - if self._embedding_model_name is None: - self._embedding_model_name = self.DEFAULT_EMBEDDING_MODEL - return self._embedding_model_name - - @property - def sparse_embedding_model_name(self) -> str | None: - return self._sparse_embedding_model_name - - def set_model( - self, - embedding_model_name: str, - max_length: int | None = None, - cache_dir: str | None = None, - threads: int | None = None, - providers: Sequence["OnnxProvider"] | None = None, - cuda: bool = False, - device_ids: list[int] | None = None, - lazy_load: bool = False, - **kwargs: Any, - ) -> None: - """ - Set embedding model to use for encoding documents and queries. - - Args: - embedding_model_name: One of the supported embedding models. See `SUPPORTED_EMBEDDING_MODELS` for details. - max_length (int, optional): Deprecated. Defaults to None. - cache_dir (str, optional): The path to the cache directory. - Can be set using the `FASTEMBED_CACHE_PATH` env variable. - Defaults to `fastembed_cache` in the system's temp directory. - threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None. - providers: The list of onnx providers (with or without options) to use. Defaults to None. - Example configuration: - https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#configuration-options - cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers` - Defaults to False. - device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in - workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None. - lazy_load (bool, optional): Whether to load the model during class initialization or on demand. - Should be set to True when using multiple-gpu and parallel encoding. Defaults to False. - Raises: - ValueError: If embedding model is not supported. - ImportError: If fastembed is not installed. - - Returns: - None - """ - if max_length is not None: - show_warning( - message="max_length parameter is deprecated and will be removed in the future. It's not used by fastembed models.", - category=DeprecationWarning, - stacklevel=3, - ) - self._get_or_init_model( - model_name=embedding_model_name, - cache_dir=cache_dir, - threads=threads, - providers=providers, - cuda=cuda, - device_ids=device_ids, - lazy_load=lazy_load, - deprecated=True, - **kwargs, - ) - self._embedding_model_name = embedding_model_name - - def set_sparse_model( - self, - embedding_model_name: str | None, - cache_dir: str | None = None, - threads: int | None = None, - providers: Sequence["OnnxProvider"] | None = None, - cuda: bool = False, - device_ids: list[int] | None = None, - lazy_load: bool = False, - **kwargs: Any, - ) -> None: - """ - Set sparse embedding model to use for hybrid search over documents in combination with dense embeddings. - - Args: - embedding_model_name: One of the supported sparse embedding models. See `SUPPORTED_SPARSE_EMBEDDING_MODELS` for details. - If None, sparse embeddings will not be used. - cache_dir (str, optional): The path to the cache directory. - Can be set using the `FASTEMBED_CACHE_PATH` env variable. - Defaults to `fastembed_cache` in the system's temp directory. - threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None. - providers: The list of onnx providers (with or without options) to use. Defaults to None. - Example configuration: - https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#configuration-options - cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers` - Defaults to False. - device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in - workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None. - lazy_load (bool, optional): Whether to load the model during class initialization or on demand. - Should be set to True when using multiple-gpu and parallel encoding. Defaults to False. - Raises: - ValueError: If embedding model is not supported. - ImportError: If fastembed is not installed. - - Returns: - None - """ - if embedding_model_name is not None: - self._get_or_init_sparse_model( - model_name=embedding_model_name, - cache_dir=cache_dir, - threads=threads, - providers=providers, - cuda=cuda, - device_ids=device_ids, - lazy_load=lazy_load, - deprecated=True, - **kwargs, - ) - self._sparse_embedding_model_name = embedding_model_name - @classmethod def _get_model_params(cls, model_name: str) -> tuple[int, models.Distance]: FastEmbedMisc.import_fastembed() @@ -229,518 +97,20 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase): ) raise ValueError(f"Unsupported embedding model: {model_name}") - def _get_or_init_model( - self, - model_name: str, - cache_dir: str | None = None, - threads: int | None = None, - providers: Sequence["OnnxProvider"] | None = None, - deprecated: bool = False, - **kwargs: Any, - ) -> "TextEmbedding": - FastEmbedMisc.import_fastembed() - assert isinstance(self._model_embedder.embedder, Embedder) - return self._model_embedder.embedder.get_or_init_model( - model_name=model_name, - cache_dir=cache_dir, - threads=threads, - providers=providers, - deprecated=deprecated, - **kwargs, - ) - - def _get_or_init_sparse_model( - self, - model_name: str, - cache_dir: str | None = None, - threads: int | None = None, - providers: Sequence["OnnxProvider"] | None = None, - deprecated: bool = False, - **kwargs: Any, - ) -> "SparseTextEmbedding": - FastEmbedMisc.import_fastembed() - assert isinstance(self._model_embedder.embedder, Embedder) - return self._model_embedder.embedder.get_or_init_sparse_model( - model_name=model_name, - cache_dir=cache_dir, - threads=threads, - providers=providers, - deprecated=deprecated, - **kwargs, - ) - - def _embed_documents( - self, - documents: Iterable[str], - embedding_model_name: str = DEFAULT_EMBEDDING_MODEL, - batch_size: int = 32, - embed_type: str = "default", - parallel: int | None = None, - ) -> Iterable[tuple[str, list[float]]]: - embedding_model = self._get_or_init_model(model_name=embedding_model_name, deprecated=True) - (documents_a, documents_b) = tee(documents, 2) - if embed_type == "passage": - vectors_iter = embedding_model.passage_embed( - documents_a, batch_size=batch_size, parallel=parallel - ) - elif embed_type == "query": - vectors_iter = ( - list(embedding_model.query_embed(query=query))[0] for query in documents_a - ) - elif embed_type == "default": - vectors_iter = embedding_model.embed( - documents_a, batch_size=batch_size, parallel=parallel - ) - else: - raise ValueError(f"Unknown embed type: {embed_type}") - for vector, doc in zip(vectors_iter, documents_b): - yield (doc, vector.tolist()) - - def _sparse_embed_documents( - self, - documents: Iterable[str], - embedding_model_name: str = DEFAULT_EMBEDDING_MODEL, - batch_size: int = 32, - parallel: int | None = None, - ) -> Iterable[types.SparseVector]: - sparse_embedding_model = self._get_or_init_sparse_model( - model_name=embedding_model_name, deprecated=True - ) - vectors_iter = sparse_embedding_model.embed( - documents, batch_size=batch_size, parallel=parallel - ) - for sparse_vector in vectors_iter: - yield types.SparseVector( - indices=sparse_vector.indices.tolist(), values=sparse_vector.values.tolist() - ) - - def get_vector_field_name(self) -> str: - """ - Returns name of the vector field in qdrant collection, used by current fastembed model. - Returns: - Name of the vector field. - """ - model_name = self.embedding_model_name.split("/")[-1].lower() - return f"fast-{model_name}" - - def get_sparse_vector_field_name(self) -> str | None: - """ - Returns name of the vector field in qdrant collection, used by current fastembed model. - Returns: - Name of the vector field. - """ - if self.sparse_embedding_model_name is not None: - model_name = self.sparse_embedding_model_name.split("/")[-1].lower() - return f"fast-sparse-{model_name}" - return None - - def _scored_points_to_query_responses( - self, scored_points: list[types.ScoredPoint] - ) -> list[QueryResponse]: - response = [] - vector_field_name = self.get_vector_field_name() - sparse_vector_field_name = self.get_sparse_vector_field_name() - for scored_point in scored_points: - embedding = ( - scored_point.vector.get(vector_field_name, None) - if isinstance(scored_point.vector, dict) - else None - ) - sparse_embedding = None - if sparse_vector_field_name is not None: - sparse_embedding = ( - scored_point.vector.get(sparse_vector_field_name, None) - if isinstance(scored_point.vector, dict) - else None - ) - response.append( - QueryResponse( - id=scored_point.id, - embedding=embedding, - sparse_embedding=sparse_embedding, - metadata=scored_point.payload, - document=scored_point.payload.get("document", ""), - score=scored_point.score, - ) - ) - return response - - def _points_iterator( - self, - ids: Iterable[models.ExtendedPointId] | None, - metadata: Iterable[dict[str, Any]] | None, - encoded_docs: Iterable[tuple[str, list[float]]], - ids_accumulator: list, - sparse_vectors: Iterable[types.SparseVector] | None = None, - ) -> Iterable[models.PointStruct]: - if ids is None: - ids = iter(lambda: uuid.uuid4().hex, None) - if metadata is None: - metadata = iter(lambda: {}, None) - if sparse_vectors is None: - sparse_vectors = iter(lambda: None, True) - vector_name = self.get_vector_field_name() - sparse_vector_name = self.get_sparse_vector_field_name() - for idx, meta, (doc, vector), sparse_vector in zip( - ids, metadata, encoded_docs, sparse_vectors - ): - ids_accumulator.append(idx) - payload = {"document": doc, **meta} - point_vector: dict[str, models.Vector] = {vector_name: vector} - if sparse_vector_name is not None and sparse_vector is not None: - point_vector[sparse_vector_name] = sparse_vector - yield models.PointStruct(id=idx, payload=payload, vector=point_vector) - - def _validate_collection_info(self, collection_info: models.CollectionInfo) -> None: - (embeddings_size, distance) = self._get_model_params(model_name=self.embedding_model_name) - vector_field_name = self.get_vector_field_name() - assert isinstance( - collection_info.config.params.vectors, dict - ), f"Collection have incompatible vector params: {collection_info.config.params.vectors}" - assert ( - vector_field_name in collection_info.config.params.vectors - ), f"Collection have incompatible vector params: {collection_info.config.params.vectors}, expected {vector_field_name}" - vector_params = collection_info.config.params.vectors[vector_field_name] - assert ( - embeddings_size == vector_params.size - ), f"Embedding size mismatch: {embeddings_size} != {vector_params.size}" - assert ( - distance == vector_params.distance - ), f"Distance mismatch: {distance} != {vector_params.distance}" - sparse_vector_field_name = self.get_sparse_vector_field_name() - if sparse_vector_field_name is not None: - assert ( - sparse_vector_field_name in collection_info.config.params.sparse_vectors - ), f"Collection have incompatible vector params: {collection_info.config.params.vectors}" - if self.sparse_embedding_model_name in IDF_EMBEDDING_MODELS: - modifier = collection_info.config.params.sparse_vectors[ - sparse_vector_field_name - ].modifier - assert ( - modifier == models.Modifier.IDF - ), f"{self.sparse_embedding_model_name} requires modifier IDF, current modifier is {modifier}" - - def get_embedding_size(self, model_name: str | None = None) -> int: + def get_embedding_size(self, model_name: str) -> int: """Get the size of the embeddings produced by the specified model. Args: - model_name: optional, the name of the model to get the embedding size for. If None, the default model will - be used. - + model_name: the name of the model to get the embedding size for. Returns: int: the size of the embeddings produced by the model. Raises: ValueError: If sparse model name is passed or model is not found in the supported models. """ - model_name = model_name or self.embedding_model_name (embeddings_size, _) = self._get_model_params(model_name=model_name) return embeddings_size - def get_fastembed_vector_params( - self, - on_disk: bool | None = None, - quantization_config: models.QuantizationConfig | None = None, - hnsw_config: models.HnswConfigDiff | None = None, - ) -> dict[str, models.VectorParams]: - """ - Generates vector configuration, compatible with fastembed models. - - Args: - on_disk: if True, vectors will be stored on disk. If None, default value will be used. - quantization_config: Quantization configuration. If None, quantization will be disabled. - hnsw_config: HNSW configuration. If None, default configuration will be used. - - Returns: - Configuration for `vectors_config` argument in `create_collection` method. - """ - vector_field_name = self.get_vector_field_name() - (embeddings_size, distance) = self._get_model_params(model_name=self.embedding_model_name) - return { - vector_field_name: models.VectorParams( - size=embeddings_size, - distance=distance, - on_disk=on_disk, - quantization_config=quantization_config, - hnsw_config=hnsw_config, - ) - } - - def get_fastembed_sparse_vector_params( - self, on_disk: bool | None = None, modifier: models.Modifier | None = None - ) -> dict[str, models.SparseVectorParams] | None: - """ - Generates vector configuration, compatible with fastembed sparse models. - - Args: - on_disk: if True, vectors will be stored on disk. If None, default value will be used. - modifier: Sparse vector queries modifier. E.g. Modifier.IDF for idf-based rescoring. Default: None. - Returns: - Configuration for `vectors_config` argument in `create_collection` method. - """ - vector_field_name = self.get_sparse_vector_field_name() - if self.sparse_embedding_model_name in IDF_EMBEDDING_MODELS: - modifier = models.Modifier.IDF if modifier is None else modifier - if vector_field_name is None: - return None - return { - vector_field_name: models.SparseVectorParams( - index=models.SparseIndexParams(on_disk=on_disk), modifier=modifier - ) - } - - async def add( - self, - collection_name: str, - documents: Iterable[str], - metadata: Iterable[dict[str, Any]] | None = None, - ids: Iterable[models.ExtendedPointId] | None = None, - batch_size: int = 32, - parallel: int | None = None, - **kwargs: Any, - ) -> list[str | int]: - """ - Adds text documents into qdrant collection. - If collection does not exist, it will be created with default parameters. - Metadata in combination with documents will be added as payload. - Documents will be embedded using the specified embedding model. - - If you want to use your own vectors, use `upsert` method instead. - - Args: - collection_name (str): - Name of the collection to add documents to. - documents (Iterable[str]): - List of documents to embed and add to the collection. - metadata (Iterable[dict[str, Any]], optional): - List of metadata dicts. Defaults to None. - ids (Iterable[models.ExtendedPointId], optional): - List of ids to assign to documents. - If not specified, UUIDs will be generated. Defaults to None. - batch_size (int, optional): - How many documents to embed and upload in single request. Defaults to 32. - parallel (Optional[int], optional): - How many parallel workers to use for embedding. Defaults to None. - If number is specified, data-parallel process will be used. - - Raises: - ImportError: If fastembed is not installed. - - Returns: - List of IDs of added documents. If no ids provided, UUIDs will be randomly generated on client side. - - """ - show_warning_once( - "`add` method has been deprecated and will be removed in 1.17. Instead, inference can be done internally within regular methods like `upsert` by wrapping data into `models.Document` or `models.Image`." - ) - encoded_docs = self._embed_documents( - documents=documents, - embedding_model_name=self.embedding_model_name, - batch_size=batch_size, - embed_type="passage", - parallel=parallel, - ) - encoded_sparse_docs = None - if self.sparse_embedding_model_name is not None: - encoded_sparse_docs = self._sparse_embed_documents( - documents=documents, - embedding_model_name=self.sparse_embedding_model_name, - batch_size=batch_size, - parallel=parallel, - ) - try: - collection_info = await self.get_collection(collection_name=collection_name) - except Exception: - await self.create_collection( - collection_name=collection_name, - vectors_config=self.get_fastembed_vector_params(), - sparse_vectors_config=self.get_fastembed_sparse_vector_params(), - ) - collection_info = await self.get_collection(collection_name=collection_name) - self._validate_collection_info(collection_info) - inserted_ids: list = [] - points = self._points_iterator( - ids=ids, - metadata=metadata, - encoded_docs=encoded_docs, - ids_accumulator=inserted_ids, - sparse_vectors=encoded_sparse_docs, - ) - self.upload_points( - collection_name=collection_name, - points=points, - wait=True, - parallel=parallel or 1, - batch_size=batch_size, - **kwargs, - ) - return inserted_ids - - async def query( - self, - collection_name: str, - query_text: str, - query_filter: models.Filter | None = None, - limit: int = 10, - **kwargs: Any, - ) -> list[QueryResponse]: - """ - Search for documents in a collection. - This method automatically embeds the query text using the specified embedding model. - If you want to use your own query vector, use `search` method instead. - - Args: - collection_name: Collection to search in - query_text: - Text to search for. This text will be embedded using the specified embedding model. - And then used as a query vector. - query_filter: - - Exclude vectors which doesn't fit given conditions. - - If `None` - search among all vectors - limit: How many results return - **kwargs: Additional search parameters. See `qdrant_client.models.QueryRequest` for details. - - Returns: - list[types.ScoredPoint]: List of scored points. - - """ - show_warning_once( - "`query` method has been deprecated and will be removed in 1.17. Instead, inference can be done internally within regular methods like `query_points` by wrapping data into `models.Document` or `models.Image`." - ) - embedding_model_inst = self._get_or_init_model( - model_name=self.embedding_model_name, deprecated=True - ) - embeddings = list(embedding_model_inst.query_embed(query=query_text)) - query_vector = embeddings[0].tolist() - if self.sparse_embedding_model_name is None: - return self._scored_points_to_query_responses( - ( - await self.query_points( - collection_name=collection_name, - query=query_vector, - using=self.get_vector_field_name(), - query_filter=query_filter, - limit=limit, - with_payload=True, - **kwargs, - ) - ).points - ) - sparse_embedding_model_inst = self._get_or_init_sparse_model( - model_name=self.sparse_embedding_model_name, deprecated=True - ) - sparse_vector = list(sparse_embedding_model_inst.query_embed(query=query_text))[0] - sparse_query_vector = models.SparseVector( - indices=sparse_vector.indices.tolist(), values=sparse_vector.values.tolist() - ) - dense_request = models.QueryRequest( - query=query_vector, - using=self.get_vector_field_name(), - filter=query_filter, - limit=limit, - with_payload=True, - **kwargs, - ) - sparse_request = models.QueryRequest( - query=sparse_query_vector, - using=self.get_sparse_vector_field_name(), - filter=query_filter, - limit=limit, - with_payload=True, - **kwargs, - ) - (dense_request_response, sparse_request_response) = await self.query_batch_points( - collection_name=collection_name, requests=[dense_request, sparse_request] - ) - return self._scored_points_to_query_responses( - reciprocal_rank_fusion( - [dense_request_response.points, sparse_request_response.points], limit=limit - ) - ) - - async def query_batch( - self, - collection_name: str, - query_texts: list[str], - query_filter: models.Filter | None = None, - limit: int = 10, - **kwargs: Any, - ) -> list[list[QueryResponse]]: - """ - Search for documents in a collection with batched query. - This method automatically embeds the query text using the specified embedding model. - - Args: - collection_name: Collection to search in - query_texts: - A list of texts to search for. Each text will be embedded using the specified embedding model. - And then used as a query vector for a separate search requests. - query_filter: - - Exclude vectors which doesn't fit given conditions. - - If `None` - search among all vectors - This filter will be applied to all search requests. - limit: How many results return - **kwargs: Additional search parameters. See `qdrant_client.models.QueryRequest` for details. - - Returns: - list[list[QueryResponse]]: List of lists of responses for each query text. - - """ - show_warning_once( - "`query_batch` method has been deprecated and will be removed in 1.17. Instead, inference can be done internally within regular methods like `query_batch_points` by wrapping data into `models.Document` or `models.Image`." - ) - embedding_model_inst = self._get_or_init_model( - model_name=self.embedding_model_name, deprecated=True - ) - query_vectors = list(embedding_model_inst.query_embed(query=query_texts)) - requests = [] - for vector in query_vectors: - request = models.QueryRequest( - query=vector.tolist(), - using=self.get_vector_field_name(), - filter=query_filter, - limit=limit, - with_payload=True, - **kwargs, - ) - requests.append(request) - if self.sparse_embedding_model_name is None: - responses = await self.query_batch_points( - collection_name=collection_name, requests=requests - ) - return [ - self._scored_points_to_query_responses(response.points) for response in responses - ] - sparse_embedding_model_inst = self._get_or_init_sparse_model( - model_name=self.sparse_embedding_model_name, deprecated=True - ) - sparse_query_vectors = [ - models.SparseVector( - indices=sparse_vector.indices.tolist(), values=sparse_vector.values.tolist() - ) - for sparse_vector in sparse_embedding_model_inst.embed(documents=query_texts) - ] - for sparse_vector in sparse_query_vectors: - request = models.QueryRequest( - using=self.get_sparse_vector_field_name(), - query=sparse_vector, - filter=query_filter, - limit=limit, - with_payload=True, - **kwargs, - ) - requests.append(request) - responses = await self.query_batch_points( - collection_name=collection_name, requests=requests - ) - dense_responses = responses[: len(query_texts)] - sparse_responses = responses[len(query_texts) :] - responses = [ - reciprocal_rank_fusion([dense_response.points, sparse_response.points], limit=limit) - for (dense_response, sparse_response) in zip(dense_responses, sparse_responses) - ] - return [self._scored_points_to_query_responses(response) for response in responses] - @classmethod def _resolve_query( cls, diff --git a/qdrant_client/qdrant_fastembed.py b/qdrant_client/qdrant_fastembed.py index 8e83293c..c9b0a466 100644 --- a/qdrant_client/qdrant_fastembed.py +++ b/qdrant_client/qdrant_fastembed.py @@ -1,5 +1,3 @@ -import uuid -from itertools import tee from typing import Any, Iterable, Sequence, get_args from copy import deepcopy @@ -7,38 +5,22 @@ import numpy as np from pydantic import BaseModel from qdrant_client import grpc -from qdrant_client.common.client_warnings import show_warning, show_warning_once from qdrant_client.client_base import QdrantBase -from qdrant_client.embed.embedder import Embedder from qdrant_client.embed.model_embedder import ModelEmbedder from qdrant_client.http import models from qdrant_client.conversions import common_types as types from qdrant_client.conversions.conversion import GrpcToRest from qdrant_client.embed.common import INFERENCE_OBJECT_TYPES from qdrant_client.embed.schema_parser import ModelSchemaParser -from qdrant_client.hybrid.fusion import reciprocal_rank_fusion -from qdrant_client.fastembed_common import FastEmbedMisc, OnnxProvider - -# region imports used in deprecated methods -from qdrant_client.fastembed_common import ( - QueryResponse, - TextEmbedding, - SparseTextEmbedding, - IDF_EMBEDDING_MODELS, -) -# endregion +from qdrant_client.fastembed_common import FastEmbedMisc class QdrantFastembedMixin(QdrantBase): - DEFAULT_EMBEDDING_MODEL = "BAAI/bge-small-en" DEFAULT_BATCH_SIZE = 8 _FASTEMBED_INSTALLED: bool def __init__(self, parser: ModelSchemaParser, is_local_mode: bool): self.__class__._FASTEMBED_INSTALLED = FastEmbedMisc.is_installed() - self._embedding_model_name: str | None = None - self._sparse_embedding_model_name: str | None = None - self._model_embedder = ModelEmbedder(parser=parser, is_local_mode=is_local_mode) super().__init__() @@ -87,127 +69,6 @@ class QdrantFastembedMixin(QdrantBase): """ return FastEmbedMisc.list_sparse_models() - @property - def embedding_model_name(self) -> str: - if self._embedding_model_name is None: - self._embedding_model_name = self.DEFAULT_EMBEDDING_MODEL - return self._embedding_model_name - - @property - def sparse_embedding_model_name(self) -> str | None: - return self._sparse_embedding_model_name - - def set_model( - self, - embedding_model_name: str, - max_length: int | None = None, - cache_dir: str | None = None, - threads: int | None = None, - providers: Sequence["OnnxProvider"] | None = None, - cuda: bool = False, - device_ids: list[int] | None = None, - lazy_load: bool = False, - **kwargs: Any, - ) -> None: - """ - Set embedding model to use for encoding documents and queries. - - Args: - embedding_model_name: One of the supported embedding models. See `SUPPORTED_EMBEDDING_MODELS` for details. - max_length (int, optional): Deprecated. Defaults to None. - cache_dir (str, optional): The path to the cache directory. - Can be set using the `FASTEMBED_CACHE_PATH` env variable. - Defaults to `fastembed_cache` in the system's temp directory. - threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None. - providers: The list of onnx providers (with or without options) to use. Defaults to None. - Example configuration: - https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#configuration-options - cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers` - Defaults to False. - device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in - workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None. - lazy_load (bool, optional): Whether to load the model during class initialization or on demand. - Should be set to True when using multiple-gpu and parallel encoding. Defaults to False. - Raises: - ValueError: If embedding model is not supported. - ImportError: If fastembed is not installed. - - Returns: - None - """ - - if max_length is not None: - show_warning( - message="max_length parameter is deprecated and will be removed in the future. " - "It's not used by fastembed models.", - category=DeprecationWarning, - stacklevel=3, - ) - - self._get_or_init_model( - model_name=embedding_model_name, - cache_dir=cache_dir, - threads=threads, - providers=providers, - cuda=cuda, - device_ids=device_ids, - lazy_load=lazy_load, - deprecated=True, - **kwargs, - ) - self._embedding_model_name = embedding_model_name - - def set_sparse_model( - self, - embedding_model_name: str | None, - cache_dir: str | None = None, - threads: int | None = None, - providers: Sequence["OnnxProvider"] | None = None, - cuda: bool = False, - device_ids: list[int] | None = None, - lazy_load: bool = False, - **kwargs: Any, - ) -> None: - """ - Set sparse embedding model to use for hybrid search over documents in combination with dense embeddings. - - Args: - embedding_model_name: One of the supported sparse embedding models. See `SUPPORTED_SPARSE_EMBEDDING_MODELS` for details. - If None, sparse embeddings will not be used. - cache_dir (str, optional): The path to the cache directory. - Can be set using the `FASTEMBED_CACHE_PATH` env variable. - Defaults to `fastembed_cache` in the system's temp directory. - threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None. - providers: The list of onnx providers (with or without options) to use. Defaults to None. - Example configuration: - https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#configuration-options - cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers` - Defaults to False. - device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in - workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None. - lazy_load (bool, optional): Whether to load the model during class initialization or on demand. - Should be set to True when using multiple-gpu and parallel encoding. Defaults to False. - Raises: - ValueError: If embedding model is not supported. - ImportError: If fastembed is not installed. - - Returns: - None - """ - if embedding_model_name is not None: - self._get_or_init_sparse_model( - model_name=embedding_model_name, - cache_dir=cache_dir, - threads=threads, - providers=providers, - cuda=cuda, - device_ids=device_ids, - lazy_load=lazy_load, - deprecated=True, - **kwargs, - ) - self._sparse_embedding_model_name = embedding_model_name - @classmethod def _get_model_params(cls, model_name: str) -> tuple[int, models.Distance]: FastEmbedMisc.import_fastembed() @@ -228,576 +89,23 @@ class QdrantFastembedMixin(QdrantBase): raise ValueError(f"Unsupported embedding model: {model_name}") - def _get_or_init_model( - self, - model_name: str, - cache_dir: str | None = None, - threads: int | None = None, - providers: Sequence["OnnxProvider"] | None = None, - deprecated: bool = False, - **kwargs: Any, - ) -> "TextEmbedding": - FastEmbedMisc.import_fastembed() - - assert isinstance(self._model_embedder.embedder, Embedder) - return self._model_embedder.embedder.get_or_init_model( - model_name=model_name, - cache_dir=cache_dir, - threads=threads, - providers=providers, - deprecated=deprecated, - **kwargs, - ) - - def _get_or_init_sparse_model( - self, - model_name: str, - cache_dir: str | None = None, - threads: int | None = None, - providers: Sequence["OnnxProvider"] | None = None, - deprecated: bool = False, - **kwargs: Any, - ) -> "SparseTextEmbedding": - FastEmbedMisc.import_fastembed() - assert isinstance(self._model_embedder.embedder, Embedder) - return self._model_embedder.embedder.get_or_init_sparse_model( - model_name=model_name, - cache_dir=cache_dir, - threads=threads, - providers=providers, - deprecated=deprecated, - **kwargs, - ) - - def _embed_documents( - self, - documents: Iterable[str], - embedding_model_name: str = DEFAULT_EMBEDDING_MODEL, - batch_size: int = 32, - embed_type: str = "default", - parallel: int | None = None, - ) -> Iterable[tuple[str, list[float]]]: - embedding_model = self._get_or_init_model(model_name=embedding_model_name, deprecated=True) - documents_a, documents_b = tee(documents, 2) - if embed_type == "passage": - vectors_iter = embedding_model.passage_embed( - documents_a, batch_size=batch_size, parallel=parallel - ) - elif embed_type == "query": - vectors_iter = ( - list(embedding_model.query_embed(query=query))[0] for query in documents_a - ) - elif embed_type == "default": - vectors_iter = embedding_model.embed( - documents_a, batch_size=batch_size, parallel=parallel - ) - else: - raise ValueError(f"Unknown embed type: {embed_type}") - - for vector, doc in zip(vectors_iter, documents_b): - yield doc, vector.tolist() - - def _sparse_embed_documents( - self, - documents: Iterable[str], - embedding_model_name: str = DEFAULT_EMBEDDING_MODEL, - batch_size: int = 32, - parallel: int | None = None, - ) -> Iterable[types.SparseVector]: - sparse_embedding_model = self._get_or_init_sparse_model( - model_name=embedding_model_name, deprecated=True - ) - - vectors_iter = sparse_embedding_model.embed( - documents, batch_size=batch_size, parallel=parallel - ) - - for sparse_vector in vectors_iter: - yield types.SparseVector( - indices=sparse_vector.indices.tolist(), - values=sparse_vector.values.tolist(), - ) - - def get_vector_field_name(self) -> str: - """ - Returns name of the vector field in qdrant collection, used by current fastembed model. - Returns: - Name of the vector field. - """ - model_name = self.embedding_model_name.split("/")[-1].lower() - return f"fast-{model_name}" - - def get_sparse_vector_field_name(self) -> str | None: - """ - Returns name of the vector field in qdrant collection, used by current fastembed model. - Returns: - Name of the vector field. - """ - if self.sparse_embedding_model_name is not None: - model_name = self.sparse_embedding_model_name.split("/")[-1].lower() - return f"fast-sparse-{model_name}" - return None - - def _scored_points_to_query_responses( - self, - scored_points: list[types.ScoredPoint], - ) -> list[QueryResponse]: - response = [] - vector_field_name = self.get_vector_field_name() - sparse_vector_field_name = self.get_sparse_vector_field_name() - - for scored_point in scored_points: - embedding = ( - scored_point.vector.get(vector_field_name, None) - if isinstance(scored_point.vector, dict) - else None - ) - sparse_embedding = None - if sparse_vector_field_name is not None: - sparse_embedding = ( - scored_point.vector.get(sparse_vector_field_name, None) - if isinstance(scored_point.vector, dict) - else None - ) - - response.append( - QueryResponse( - id=scored_point.id, - embedding=embedding, - sparse_embedding=sparse_embedding, - metadata=scored_point.payload, - document=scored_point.payload.get("document", ""), - score=scored_point.score, - ) - ) - return response - - def _points_iterator( - self, - ids: Iterable[models.ExtendedPointId] | None, - metadata: Iterable[dict[str, Any]] | None, - encoded_docs: Iterable[tuple[str, list[float]]], - ids_accumulator: list, - sparse_vectors: Iterable[types.SparseVector] | None = None, - ) -> Iterable[models.PointStruct]: - if ids is None: - ids = iter(lambda: uuid.uuid4().hex, None) - - if metadata is None: - metadata = iter(lambda: {}, None) - - if sparse_vectors is None: - sparse_vectors = iter(lambda: None, True) - - vector_name = self.get_vector_field_name() - sparse_vector_name = self.get_sparse_vector_field_name() - - for idx, meta, (doc, vector), sparse_vector in zip( - ids, metadata, encoded_docs, sparse_vectors - ): - ids_accumulator.append(idx) - payload = {"document": doc, **meta} - point_vector: dict[str, models.Vector] = {vector_name: vector} - if sparse_vector_name is not None and sparse_vector is not None: - point_vector[sparse_vector_name] = sparse_vector - yield models.PointStruct(id=idx, payload=payload, vector=point_vector) - - def _validate_collection_info(self, collection_info: models.CollectionInfo) -> None: - embeddings_size, distance = self._get_model_params(model_name=self.embedding_model_name) - vector_field_name = self.get_vector_field_name() - - # Check if collection has compatible vector params - assert isinstance( - collection_info.config.params.vectors, dict - ), f"Collection have incompatible vector params: {collection_info.config.params.vectors}" - - assert ( - vector_field_name in collection_info.config.params.vectors - ), f"Collection have incompatible vector params: {collection_info.config.params.vectors}, expected {vector_field_name}" - - vector_params = collection_info.config.params.vectors[vector_field_name] - - assert ( - embeddings_size == vector_params.size - ), f"Embedding size mismatch: {embeddings_size} != {vector_params.size}" - - assert ( - distance == vector_params.distance - ), f"Distance mismatch: {distance} != {vector_params.distance}" - - sparse_vector_field_name = self.get_sparse_vector_field_name() - if sparse_vector_field_name is not None: - assert ( - sparse_vector_field_name in collection_info.config.params.sparse_vectors - ), f"Collection have incompatible vector params: {collection_info.config.params.vectors}" - if self.sparse_embedding_model_name in IDF_EMBEDDING_MODELS: - modifier = collection_info.config.params.sparse_vectors[ - sparse_vector_field_name - ].modifier - assert ( - modifier == models.Modifier.IDF - ), f"{self.sparse_embedding_model_name} requires modifier IDF, current modifier is {modifier}" - def get_embedding_size( self, - model_name: str | None = None, + model_name: str, ) -> int: """Get the size of the embeddings produced by the specified model. Args: - model_name: optional, the name of the model to get the embedding size for. If None, the default model will - be used. - + model_name: the name of the model to get the embedding size for. Returns: int: the size of the embeddings produced by the model. Raises: ValueError: If sparse model name is passed or model is not found in the supported models. """ - model_name = model_name or self.embedding_model_name embeddings_size, _ = self._get_model_params(model_name=model_name) return embeddings_size - def get_fastembed_vector_params( - self, - on_disk: bool | None = None, - quantization_config: models.QuantizationConfig | None = None, - hnsw_config: models.HnswConfigDiff | None = None, - ) -> dict[str, models.VectorParams]: - """ - Generates vector configuration, compatible with fastembed models. - - Args: - on_disk: if True, vectors will be stored on disk. If None, default value will be used. - quantization_config: Quantization configuration. If None, quantization will be disabled. - hnsw_config: HNSW configuration. If None, default configuration will be used. - - Returns: - Configuration for `vectors_config` argument in `create_collection` method. - """ - vector_field_name = self.get_vector_field_name() - embeddings_size, distance = self._get_model_params(model_name=self.embedding_model_name) - return { - vector_field_name: models.VectorParams( - size=embeddings_size, - distance=distance, - on_disk=on_disk, - quantization_config=quantization_config, - hnsw_config=hnsw_config, - ) - } - - def get_fastembed_sparse_vector_params( - self, - on_disk: bool | None = None, - modifier: models.Modifier | None = None, - ) -> dict[str, models.SparseVectorParams] | None: - """ - Generates vector configuration, compatible with fastembed sparse models. - - Args: - on_disk: if True, vectors will be stored on disk. If None, default value will be used. - modifier: Sparse vector queries modifier. E.g. Modifier.IDF for idf-based rescoring. Default: None. - Returns: - Configuration for `vectors_config` argument in `create_collection` method. - """ - vector_field_name = self.get_sparse_vector_field_name() - if self.sparse_embedding_model_name in IDF_EMBEDDING_MODELS: - modifier = models.Modifier.IDF if modifier is None else modifier - - if vector_field_name is None: - return None - - return { - vector_field_name: models.SparseVectorParams( - index=models.SparseIndexParams( - on_disk=on_disk, - ), - modifier=modifier, - ) - } - - def add( - self, - collection_name: str, - documents: Iterable[str], - metadata: Iterable[dict[str, Any]] | None = None, - ids: Iterable[models.ExtendedPointId] | None = None, - batch_size: int = 32, - parallel: int | None = None, - **kwargs: Any, - ) -> list[str | int]: - """ - Adds text documents into qdrant collection. - If collection does not exist, it will be created with default parameters. - Metadata in combination with documents will be added as payload. - Documents will be embedded using the specified embedding model. - - If you want to use your own vectors, use `upsert` method instead. - - Args: - collection_name (str): - Name of the collection to add documents to. - documents (Iterable[str]): - List of documents to embed and add to the collection. - metadata (Iterable[dict[str, Any]], optional): - List of metadata dicts. Defaults to None. - ids (Iterable[models.ExtendedPointId], optional): - List of ids to assign to documents. - If not specified, UUIDs will be generated. Defaults to None. - batch_size (int, optional): - How many documents to embed and upload in single request. Defaults to 32. - parallel (Optional[int], optional): - How many parallel workers to use for embedding. Defaults to None. - If number is specified, data-parallel process will be used. - - Raises: - ImportError: If fastembed is not installed. - - Returns: - List of IDs of added documents. If no ids provided, UUIDs will be randomly generated on client side. - - """ - show_warning_once( - "`add` method has been deprecated and will be removed in 1.17. " - "Instead, inference can be done internally within regular methods like `upsert` by wrapping " - "data into `models.Document` or `models.Image`." - ) - - # check if we have fastembed installed - encoded_docs = self._embed_documents( - documents=documents, - embedding_model_name=self.embedding_model_name, - batch_size=batch_size, - embed_type="passage", - parallel=parallel, - ) - - encoded_sparse_docs = None - if self.sparse_embedding_model_name is not None: - encoded_sparse_docs = self._sparse_embed_documents( - documents=documents, - embedding_model_name=self.sparse_embedding_model_name, - batch_size=batch_size, - parallel=parallel, - ) - - # Check if collection by same name exists, if not, create it - try: - collection_info = self.get_collection(collection_name=collection_name) - except Exception: - self.create_collection( - collection_name=collection_name, - vectors_config=self.get_fastembed_vector_params(), - sparse_vectors_config=self.get_fastembed_sparse_vector_params(), - ) - collection_info = self.get_collection(collection_name=collection_name) - - self._validate_collection_info(collection_info) - - inserted_ids: list = [] - - points = self._points_iterator( - ids=ids, - metadata=metadata, - encoded_docs=encoded_docs, - ids_accumulator=inserted_ids, - sparse_vectors=encoded_sparse_docs, - ) - - self.upload_points( - collection_name=collection_name, - points=points, - wait=True, - parallel=parallel or 1, - batch_size=batch_size, - **kwargs, - ) - - return inserted_ids - - def query( - self, - collection_name: str, - query_text: str, - query_filter: models.Filter | None = None, - limit: int = 10, - **kwargs: Any, - ) -> list[QueryResponse]: - """ - Search for documents in a collection. - This method automatically embeds the query text using the specified embedding model. - If you want to use your own query vector, use `search` method instead. - - Args: - collection_name: Collection to search in - query_text: - Text to search for. This text will be embedded using the specified embedding model. - And then used as a query vector. - query_filter: - - Exclude vectors which doesn't fit given conditions. - - If `None` - search among all vectors - limit: How many results return - **kwargs: Additional search parameters. See `qdrant_client.models.QueryRequest` for details. - - Returns: - list[types.ScoredPoint]: List of scored points. - - """ - show_warning_once( - "`query` method has been deprecated and will be removed in 1.17. " - "Instead, inference can be done internally within regular methods like `query_points` by wrapping " - "data into `models.Document` or `models.Image`." - ) - embedding_model_inst = self._get_or_init_model( - model_name=self.embedding_model_name, deprecated=True - ) - embeddings = list(embedding_model_inst.query_embed(query=query_text)) - query_vector = embeddings[0].tolist() - - if self.sparse_embedding_model_name is None: - return self._scored_points_to_query_responses( - self.query_points( - collection_name=collection_name, - query=query_vector, - using=self.get_vector_field_name(), - query_filter=query_filter, - limit=limit, - with_payload=True, - **kwargs, - ).points - ) - - sparse_embedding_model_inst = self._get_or_init_sparse_model( - model_name=self.sparse_embedding_model_name, deprecated=True - ) - sparse_vector = list(sparse_embedding_model_inst.query_embed(query=query_text))[0] - sparse_query_vector = models.SparseVector( - indices=sparse_vector.indices.tolist(), - values=sparse_vector.values.tolist(), - ) - - dense_request = models.QueryRequest( - query=query_vector, - using=self.get_vector_field_name(), - filter=query_filter, - limit=limit, - with_payload=True, - **kwargs, - ) - sparse_request = models.QueryRequest( - query=sparse_query_vector, - using=self.get_sparse_vector_field_name(), - filter=query_filter, - limit=limit, - with_payload=True, - **kwargs, - ) - - dense_request_response, sparse_request_response = self.query_batch_points( - collection_name=collection_name, requests=[dense_request, sparse_request] - ) - return self._scored_points_to_query_responses( - reciprocal_rank_fusion( - [dense_request_response.points, sparse_request_response.points], limit=limit - ) - ) - - def query_batch( - self, - collection_name: str, - query_texts: list[str], - query_filter: models.Filter | None = None, - limit: int = 10, - **kwargs: Any, - ) -> list[list[QueryResponse]]: - """ - Search for documents in a collection with batched query. - This method automatically embeds the query text using the specified embedding model. - - Args: - collection_name: Collection to search in - query_texts: - A list of texts to search for. Each text will be embedded using the specified embedding model. - And then used as a query vector for a separate search requests. - query_filter: - - Exclude vectors which doesn't fit given conditions. - - If `None` - search among all vectors - This filter will be applied to all search requests. - limit: How many results return - **kwargs: Additional search parameters. See `qdrant_client.models.QueryRequest` for details. - - Returns: - list[list[QueryResponse]]: List of lists of responses for each query text. - - """ - show_warning_once( - "`query_batch` method has been deprecated and will be removed in 1.17. " - "Instead, inference can be done internally within regular methods like `query_batch_points` by wrapping " - "data into `models.Document` or `models.Image`." - ) - embedding_model_inst = self._get_or_init_model( - model_name=self.embedding_model_name, deprecated=True - ) - query_vectors = list(embedding_model_inst.query_embed(query=query_texts)) - requests = [] - for vector in query_vectors: - request = models.QueryRequest( - query=vector.tolist(), - using=self.get_vector_field_name(), - filter=query_filter, - limit=limit, - with_payload=True, - **kwargs, - ) - - requests.append(request) - - if self.sparse_embedding_model_name is None: - responses = self.query_batch_points( - collection_name=collection_name, - requests=requests, - ) - return [ - self._scored_points_to_query_responses(response.points) for response in responses - ] - - sparse_embedding_model_inst = self._get_or_init_sparse_model( - model_name=self.sparse_embedding_model_name, deprecated=True - ) - sparse_query_vectors = [ - models.SparseVector( - indices=sparse_vector.indices.tolist(), - values=sparse_vector.values.tolist(), - ) - for sparse_vector in sparse_embedding_model_inst.embed(documents=query_texts) - ] - for sparse_vector in sparse_query_vectors: - request = models.QueryRequest( - using=self.get_sparse_vector_field_name(), - query=sparse_vector, - filter=query_filter, - limit=limit, - with_payload=True, - **kwargs, - ) - - requests.append(request) - - responses = self.query_batch_points( - collection_name=collection_name, - requests=requests, - ) - - dense_responses = responses[: len(query_texts)] - sparse_responses = responses[len(query_texts) :] - responses = [ - reciprocal_rank_fusion([dense_response.points, sparse_response.points], limit=limit) - for dense_response, sparse_response in zip(dense_responses, sparse_responses) - ] - - return [self._scored_points_to_query_responses(response) for response in responses] - @classmethod def _resolve_query( cls, diff --git a/tests/embed_tests/test_local_inference.py b/tests/embed_tests/test_local_inference.py index b4af9ad4..6e14db4e 100644 --- a/tests/embed_tests/test_local_inference.py +++ b/tests/embed_tests/test_local_inference.py @@ -12,12 +12,12 @@ import pytest import qdrant_client.embed.embedder from qdrant_client import QdrantClient, models from qdrant_client.client_base import QdrantBase -from qdrant_client.qdrant_fastembed import IDF_EMBEDDING_MODELS from qdrant_client.fastembed_common import ( TextEmbedding, SparseTextEmbedding, LateInteractionTextEmbedding, ImageEmbedding, + IDF_EMBEDDING_MODELS, ) COLLECTION_NAME = "inference_collection" @@ -262,7 +262,7 @@ def test_upsert(cached_embeddings): def test_upload(cached_embeddings): - def recreate_collection(client, collection_name): + def recreate_collection(client: QdrantBase, collection_name: str): if client.collection_exists(collection_name): client.delete_collection(collection_name) vector_params = { diff --git a/tests/test_fastembed.py b/tests/test_fastembed.py index 7fc3f93f..9b4b3cd8 100644 --- a/tests/test_fastembed.py +++ b/tests/test_fastembed.py @@ -1,162 +1,6 @@ import pytest -from qdrant_client import QdrantClient, models - -from tests.utils import read_version - - -DOCS_EXAMPLE = { - "documents": [ - "Qdrant has Langchain integrations", - "Qdrant also has Llama Index integrations", - ], - "metadata": [{"source": "Langchain-docs"}, {"source": "LlamaIndex-docs"}], - "ids": [42, 2000], -} - - -def test_dense(): - local_client = QdrantClient(":memory:") - collection_name = "demo_collection" - docs = [ - "Qdrant has Langchain integrations", - "Qdrant also has Llama Index integrations", - ] - - if not local_client._FASTEMBED_INSTALLED: - with pytest.raises(ImportError): - local_client.add(collection_name, docs) - else: - local_client.add(collection_name=collection_name, documents=docs) - assert local_client.count(collection_name).count == 2 - - local_client.add(collection_name=collection_name, **DOCS_EXAMPLE) - assert local_client.count(collection_name).count == 4 - - id_ = DOCS_EXAMPLE["ids"][0] - record = local_client.retrieve(collection_name, ids=[id_])[0] - assert record.payload == { - "document": DOCS_EXAMPLE["documents"][0], - **DOCS_EXAMPLE["metadata"][0], - } - - search_result = local_client.query( - collection_name=collection_name, query_text="This is a query document" - ) - - assert len(search_result) > 0 - - -def test_hybrid_query(): - local_client = QdrantClient(":memory:") - collection_name = "hybrid_collection" - - if not local_client._FASTEMBED_INSTALLED: - pytest.skip("FastEmbed is not installed, skipping test") - - local_client.set_sparse_model(embedding_model_name="prithivida/Splade_PP_en_v1") - - local_client.add(collection_name=collection_name, **DOCS_EXAMPLE) - - hybrid_search_result = local_client.query( - collection_name=collection_name, query_text="This is a query document" - ) - - assert len(hybrid_search_result) > 0 - - local_client.set_sparse_model(None) - dense_search_result = local_client.query( - collection_name=collection_name, query_text="This is a query document" - ) - assert len(dense_search_result) > 0 - - assert ( - hybrid_search_result[0].score != dense_search_result[0].score - ) # hybrid search has score from fusion - - -def test_query_batch(): - local_client = QdrantClient(":memory:") - - dense_collection_name = "dense_collection" - hybrid_collection_name = "hybrid_collection" - - if not local_client._FASTEMBED_INSTALLED: - pytest.skip("FastEmbed is not installed, skipping test") - - local_client.add(collection_name=dense_collection_name, **DOCS_EXAMPLE) - query_texts = ["This is a query document", "This is another query document"] - dense_search_result = local_client.query_batch( - collection_name=dense_collection_name, query_texts=query_texts - ) - assert len(dense_search_result) == len(query_texts) - assert all(len(result) > 0 for result in dense_search_result) - - local_client.set_sparse_model(embedding_model_name="prithivida/Splade_PP_en_v1") - - local_client.add(collection_name=hybrid_collection_name, **DOCS_EXAMPLE) - - hybrid_search_result = local_client.query_batch( - collection_name=hybrid_collection_name, query_texts=query_texts - ) - - assert len(hybrid_search_result) == len(query_texts) - assert all(len(result) > 0 for result in hybrid_search_result) - - single_dense_response = next(iter(dense_search_result)) - single_hybrid_response = next(iter(hybrid_search_result)) - - assert ( - single_hybrid_response[0].score != single_dense_response[0].score - ) # hybrid search has score from fusion - - -def test_set_model(): - local_client = QdrantClient(":memory:") - collection_name = "demo_collection" - embedding_model_name = "sentence-transformers/all-MiniLM-L6-v2" - if not local_client._FASTEMBED_INSTALLED: - pytest.skip("FastEmbed is not installed, skipping test") - - local_client.set_model( - embedding_model_name=embedding_model_name, - ) - - # Check if the model is initialized & cls.embeddings_models is set with expected values - dim, dist = local_client._get_model_params(embedding_model_name) - assert dim == 384 - - # Use the initialized model to add documents with vector embeddings - local_client.add(collection_name=collection_name, **DOCS_EXAMPLE) - assert local_client.count(collection_name).count == 2 - - -def test_idf_models(): - local_client = QdrantClient(":memory:") - - if not local_client._FASTEMBED_INSTALLED: - pytest.skip("FastEmbed is not installed, skipping") - - major, minor, patch, dev = read_version() - if not dev and None not in (major, minor, patch) and (major, minor, patch) < (1, 10, 2): - pytest.skip("Works as of version 1.10.2") - - for model_name in ("Qdrant/bm25", "Qdrant/bm42-all-minilm-l6-v2-attentions"): - local_client.set_sparse_model(model_name) - collection_name = model_name.split("/")[-1].replace("-", "_") - - local_client.add(collection_name=collection_name, **DOCS_EXAMPLE) - local_client.query( - collection_name=collection_name, query_text="Qdrant and Llama Index integration" - ) - - collection_info = local_client.get_collection(collection_name=collection_name) - vector_name = local_client.get_sparse_vector_field_name() - modifier = collection_info.config.params.sparse_vectors[vector_name].modifier - assert modifier == models.Modifier.IDF - - # the only sparse model without IDF is SPLADE, however it's too large for tests, so we don't test how non-idf - # models work +from qdrant_client import QdrantClient def test_get_embedding_size(): @@ -165,8 +9,6 @@ def test_get_embedding_size(): if not local_client._FASTEMBED_INSTALLED: pytest.skip("FastEmbed is not installed, skipping test") - assert local_client.get_embedding_size() == 384 - assert local_client.get_embedding_size(model_name="BAAI/bge-base-en-v1.5") == 768 assert local_client.get_embedding_size(model_name="Qdrant/resnet50-onnx") == 2048 diff --git a/tools/async_client_generator/fastembed_generator.py b/tools/async_client_generator/fastembed_generator.py index 1173e77f..0b4713bc 100644 --- a/tools/async_client_generator/fastembed_generator.py +++ b/tools/async_client_generator/fastembed_generator.py @@ -52,15 +52,7 @@ if __name__ == "__main__": generator = FastembedGenerator( keep_sync=[ "__init__", - "set_model", - "set_sparse_model", - "get_vector_field_name", - "get_sparse_vector_field_name", "get_embedding_size", - "get_fastembed_vector_params", - "get_fastembed_sparse_vector_params", - "embedding_model_name", - "sparse_embedding_model_name", ], class_replace_map={ "QdrantBase": "AsyncQdrantBase",