diff --git a/qdrant_client/async_qdrant_fastembed.py b/qdrant_client/async_qdrant_fastembed.py index 797ead50..eed777ef 100644 --- a/qdrant_client/async_qdrant_fastembed.py +++ b/qdrant_client/async_qdrant_fastembed.py @@ -25,18 +25,12 @@ 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, - LateInteractionTextEmbedding, - ImageEmbedding, SparseTextEmbedding, - SUPPORTED_EMBEDDING_MODELS, - SUPPORTED_SPARSE_EMBEDDING_MODELS, - _LATE_INTERACTION_EMBEDDING_MODELS, - _IMAGE_EMBEDDING_MODELS, IDF_EMBEDDING_MODELS, - OnnxProvider, ) @@ -49,16 +43,56 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase): self._embedding_model_name: Optional[str] = None self._sparse_embedding_model_name: Optional[str] = None self._model_embedder = ModelEmbedder(parser=parser, **kwargs) - try: - from fastembed import SparseTextEmbedding, TextEmbedding - - assert len(SparseTextEmbedding.list_supported_models()) > 0 - assert len(TextEmbedding.list_supported_models()) > 0 - self.__class__._FASTEMBED_INSTALLED = True - except ImportError: - self.__class__._FASTEMBED_INSTALLED = False + self.__class__._FASTEMBED_INSTALLED = FastEmbedMisc.is_installed() super().__init__(**kwargs) + @classmethod + async def list_text_models(cls) -> dict[str, tuple[int, models.Distance]]: + """Lists the supported dense text models. + + Returns: + dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics. + """ + return FastEmbedMisc.list_text_models() + + @classmethod + async def list_image_models(cls) -> dict[str, tuple[int, models.Distance]]: + """Lists the supported image dense models. + + Returns: + dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics. + """ + return FastEmbedMisc.list_image_models() + + @classmethod + async def list_late_interaction_text_models(cls) -> dict[str, tuple[int, models.Distance]]: + """Lists the supported late interaction text models. + + Returns: + dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics. + """ + return FastEmbedMisc.list_late_interaction_text_models() + + @classmethod + async def list_late_interaction_multimodal_models( + cls, + ) -> dict[str, tuple[int, models.Distance]]: + """Lists the supported late interaction multimodal models. + + Returns: + dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics. + """ + return FastEmbedMisc.list_late_interaction_multimodal_models() + + @classmethod + async def list_sparse_models(cls) -> dict[str, dict[str, Any]]: + """Lists the supported sparse text models. + + Returns: + dict[str, dict[str, Any]]: A dict of model names and their descriptions. + """ + return FastEmbedMisc.list_sparse_models() + @property def embedding_model_name(self) -> str: if self._embedding_model_name is None: @@ -177,24 +211,18 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase): ) self._sparse_embedding_model_name = embedding_model_name - @classmethod - def _import_fastembed(cls) -> None: - if cls._FASTEMBED_INSTALLED: - return - raise ImportError( - "fastembed is not installed. Please install it to enable fast vector indexing with `pip install fastembed`." - ) - @classmethod def _get_model_params(cls, model_name: str) -> tuple[int, models.Distance]: - cls._import_fastembed() - if model_name in SUPPORTED_EMBEDDING_MODELS: - return SUPPORTED_EMBEDDING_MODELS[model_name] - if model_name in _LATE_INTERACTION_EMBEDDING_MODELS: - return _LATE_INTERACTION_EMBEDDING_MODELS[model_name] - if model_name in _IMAGE_EMBEDDING_MODELS: - return _IMAGE_EMBEDDING_MODELS[model_name] - if model_name in SUPPORTED_SPARSE_EMBEDDING_MODELS: + FastEmbedMisc.import_fastembed() + for descriptions in ( + FastEmbedMisc.list_text_models(), + FastEmbedMisc.list_image_models(), + FastEmbedMisc.list_late_interaction_text_models(), + FastEmbedMisc.list_late_interaction_multimodal_models(), + ): + if params := descriptions.get(model_name): + return params + if model_name in FastEmbedMisc.list_sparse_models(): raise ValueError( "Sparse embeddings do not return fixed embedding size and distance type" ) @@ -209,7 +237,7 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase): deprecated: bool = False, **kwargs: Any, ) -> "TextEmbedding": - self._import_fastembed() + FastEmbedMisc.import_fastembed() return self._model_embedder.embedder.get_or_init_model( model_name=model_name, cache_dir=cache_dir, @@ -228,7 +256,7 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase): deprecated: bool = False, **kwargs: Any, ) -> "SparseTextEmbedding": - self._import_fastembed() + FastEmbedMisc.import_fastembed() return self._model_embedder.embedder.get_or_init_sparse_model( model_name=model_name, cache_dir=cache_dir, @@ -238,40 +266,6 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase): **kwargs, ) - def _get_or_init_late_interaction_model( - self, - model_name: str, - cache_dir: Optional[str] = None, - threads: Optional[int] = None, - providers: Optional[Sequence["OnnxProvider"]] = None, - **kwargs: Any, - ) -> "LateInteractionTextEmbedding": - self._import_fastembed() - return self._model_embedder.embedder.get_or_init_late_interaction_model( - model_name=model_name, - cache_dir=cache_dir, - threads=threads, - providers=providers, - **kwargs, - ) - - def _get_or_init_image_model( - self, - model_name: str, - cache_dir: Optional[str] = None, - threads: Optional[int] = None, - providers: Optional[Sequence["OnnxProvider"]] = None, - **kwargs: Any, - ) -> "ImageEmbedding": - self._import_fastembed() - return self._model_embedder.embedder.get_or_init_image_model( - model_name=model_name, - cache_dir=cache_dir, - threads=threads, - providers=providers, - **kwargs, - ) - def _embed_documents( self, documents: Iterable[str], @@ -424,20 +418,19 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase): ), f"{self.sparse_embedding_model_name} requires modifier IDF, current modifier is {modifier}" def get_embedding_size(self, model_name: Optional[str] = None) -> int: - """ - Get the size of the embeddings produced by the specified model. + """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: optional, the name of the model to get the embedding size for. If None, the default model will + be used. 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 - if model_name in SUPPORTED_SPARSE_EMBEDDING_MODELS: - raise ValueError( - f"Sparse embeddings do not have a fixed embedding size. Current model: {model_name}" - ) (embeddings_size, _) = self._get_model_params(model_name=model_name) return embeddings_size @@ -807,6 +800,7 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase): is_query: bool = False, batch_size: Optional[int] = None, ) -> Iterable[BaseModel]: + FastEmbedMisc.import_fastembed() yield from self._model_embedder.embed_models( raw_models=raw_models, is_query=is_query, @@ -819,6 +813,7 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase): batch_size: Optional[int] = None, parallel: Optional[int] = None, ) -> Iterable[BaseModel]: + FastEmbedMisc.import_fastembed() yield from self._model_embedder.embed_models_strict( raw_models=raw_models, batch_size=batch_size or self.DEFAULT_BATCH_SIZE, diff --git a/qdrant_client/embed/embedder.py b/qdrant_client/embed/embedder.py index 60449af9..90c944c6 100644 --- a/qdrant_client/embed/embedder.py +++ b/qdrant_client/embed/embedder.py @@ -1,5 +1,6 @@ from collections import defaultdict from typing import Optional, Sequence, Any, TypeVar, Generic + from pydantic import BaseModel from qdrant_client.http import models @@ -12,11 +13,7 @@ from qdrant_client.fastembed_common import ( LateInteractionTextEmbedding, LateInteractionMultimodalEmbedding, ImageEmbedding, - SUPPORTED_EMBEDDING_MODELS, - SUPPORTED_SPARSE_EMBEDDING_MODELS, - _LATE_INTERACTION_EMBEDDING_MODELS, - _IMAGE_EMBEDDING_MODELS, - _LATE_INTERACTION_MULTIMODAL_EMBEDDING_MODELS, + FastEmbedMisc, ) @@ -57,9 +54,9 @@ class Embedder: deprecated: bool = False, **kwargs: Any, ) -> TextEmbedding: - if model_name not in SUPPORTED_EMBEDDING_MODELS: + if not FastEmbedMisc.is_supported_text_model(model_name): raise ValueError( - f"Unsupported embedding model: {model_name}. Supported models: {SUPPORTED_EMBEDDING_MODELS}" + f"Unsupported embedding model: {model_name}. Supported models: {FastEmbedMisc.list_text_models()}" ) options = { "cache_dir": cache_dir, @@ -93,9 +90,9 @@ class Embedder: deprecated: bool = False, **kwargs: Any, ) -> SparseTextEmbedding: - if model_name not in SUPPORTED_SPARSE_EMBEDDING_MODELS: + if not FastEmbedMisc.is_supported_sparse_model(model_name): raise ValueError( - f"Unsupported embedding model: {model_name}. Supported models: {SUPPORTED_SPARSE_EMBEDDING_MODELS}" + f"Unsupported embedding model: {model_name}. Supported models: {FastEmbedMisc.list_sparse_models()}" ) options = { @@ -130,9 +127,10 @@ class Embedder: device_ids: Optional[list[int]] = None, **kwargs: Any, ) -> LateInteractionTextEmbedding: - if model_name not in _LATE_INTERACTION_EMBEDDING_MODELS: + if not FastEmbedMisc.is_supported_late_interaction_text_model(model_name): raise ValueError( - f"Unsupported embedding model: {model_name}. Supported models: {_LATE_INTERACTION_EMBEDDING_MODELS}" + f"Unsupported embedding model: {model_name}. " + f"Supported models: {FastEmbedMisc.list_late_interaction_text_models()}" ) options = { "cache_dir": cache_dir, @@ -164,9 +162,10 @@ class Embedder: device_ids: Optional[list[int]] = None, **kwargs: Any, ) -> LateInteractionMultimodalEmbedding: - if model_name not in _LATE_INTERACTION_MULTIMODAL_EMBEDDING_MODELS: + if not FastEmbedMisc.is_supported_late_interaction_multimodal_model(model_name): raise ValueError( - f"Unsupported embedding model: {model_name}. Supported models: {_LATE_INTERACTION_MULTIMODAL_EMBEDDING_MODELS}" + f"Unsupported embedding model: {model_name}. " + f"Supported models: {FastEmbedMisc.list_late_interaction_multimodal_models()}" ) options = { "cache_dir": cache_dir, @@ -198,9 +197,9 @@ class Embedder: device_ids: Optional[list[int]] = None, **kwargs: Any, ) -> ImageEmbedding: - if model_name not in _IMAGE_EMBEDDING_MODELS: + if not FastEmbedMisc.is_supported_image_model(model_name): raise ValueError( - f"Unsupported embedding model: {model_name}. Supported models: {_IMAGE_EMBEDDING_MODELS}" + f"Unsupported embedding model: {model_name}. Supported models: {FastEmbedMisc.list_image_models()}" ) options = { "cache_dir": cache_dir, @@ -234,19 +233,19 @@ class Embedder: embeddings: NumericVector # define type for a static type checker if texts is not None: - if model_name in SUPPORTED_EMBEDDING_MODELS: + if FastEmbedMisc.is_supported_text_model(model_name): embeddings = self._embed_dense_text( texts, model_name, options, is_query, batch_size ) - elif model_name in SUPPORTED_SPARSE_EMBEDDING_MODELS: + elif FastEmbedMisc.is_supported_sparse_model(model_name): embeddings = self._embed_sparse_text( texts, model_name, options, is_query, batch_size ) - elif model_name in _LATE_INTERACTION_EMBEDDING_MODELS: + elif FastEmbedMisc.is_supported_late_interaction_text_model(model_name): embeddings = self._embed_late_interaction_text( texts, model_name, options, is_query, batch_size ) - elif model_name in _LATE_INTERACTION_MULTIMODAL_EMBEDDING_MODELS: + elif FastEmbedMisc.is_supported_late_interaction_multimodal_model(model_name): embeddings = self._embed_late_interaction_multimodal_text( texts, model_name, options, batch_size ) @@ -256,9 +255,9 @@ class Embedder: assert ( images is not None ) # just to satisfy mypy which can't infer it from the previous conditions - if model_name in _IMAGE_EMBEDDING_MODELS: + if FastEmbedMisc.is_supported_image_model(model_name): embeddings = self._embed_dense_image(images, model_name, options, batch_size) - elif model_name in _LATE_INTERACTION_MULTIMODAL_EMBEDDING_MODELS: + elif FastEmbedMisc.is_supported_late_interaction_multimodal_model(model_name): embeddings = self._embed_late_interaction_multimodal_image( images, model_name, options, batch_size ) @@ -278,7 +277,10 @@ class Embedder: embedding_model_inst = self.get_or_init_model(model_name=model_name, **options or {}) if not is_query: - embeddings = [embedding.tolist() for embedding in embedding_model_inst.embed(documents=texts, batch_size=batch_size)] + embeddings = [ + embedding.tolist() + for embedding in embedding_model_inst.embed(documents=texts, batch_size=batch_size) + ] else: embeddings = [ embedding.tolist() for embedding in embedding_model_inst.query_embed(query=texts) diff --git a/qdrant_client/embed/model_embedder.py b/qdrant_client/embed/model_embedder.py index 9afedf93..2ccdda49 100644 --- a/qdrant_client/embed/model_embedder.py +++ b/qdrant_client/embed/model_embedder.py @@ -13,13 +13,7 @@ from qdrant_client.embed.embedder import Embedder from qdrant_client.embed.models import NumericVector, NumericVectorStruct from qdrant_client.embed.schema_parser import ModelSchemaParser from qdrant_client.embed.utils import FieldPath -from qdrant_client.fastembed_common import ( - SUPPORTED_EMBEDDING_MODELS, - SUPPORTED_SPARSE_EMBEDDING_MODELS, - _LATE_INTERACTION_EMBEDDING_MODELS, - _IMAGE_EMBEDDING_MODELS, - _LATE_INTERACTION_MULTIMODAL_EMBEDDING_MODELS, -) +from qdrant_client.fastembed_common import FastEmbedMisc from qdrant_client.parallel_processor import ParallelWorkerPool, Worker from qdrant_client.uploader.uploader import iter_batch @@ -278,6 +272,7 @@ class ModelEmbedder: if data.model not in self._batch_accumulator: self._batch_accumulator[data.model] = [] self._batch_accumulator[data.model].append(data) + return None def _drain_accumulator( self, data: models.VectorStruct, is_query: bool, inference_batch_size: int = 8 @@ -382,12 +377,14 @@ class ModelEmbedder: return ordered_embeddings for model in self._batch_accumulator: - if model not in ( - *SUPPORTED_EMBEDDING_MODELS, - *SUPPORTED_SPARSE_EMBEDDING_MODELS, - *_LATE_INTERACTION_EMBEDDING_MODELS, - *_IMAGE_EMBEDDING_MODELS, - *_LATE_INTERACTION_MULTIMODAL_EMBEDDING_MODELS, + if not any( + ( + FastEmbedMisc.is_supported_text_model(model), + FastEmbedMisc.is_supported_sparse_model(model), + FastEmbedMisc.is_supported_late_interaction_text_model(model), + FastEmbedMisc.is_supported_image_model(model), + FastEmbedMisc.is_supported_late_interaction_multimodal_model(model), + ) ): raise ValueError(f"{model} is not among supported models") @@ -426,15 +423,17 @@ class ModelEmbedder: model_name = data.model value = data.object options = data.options - if model_name in ( - *SUPPORTED_EMBEDDING_MODELS, - *SUPPORTED_SPARSE_EMBEDDING_MODELS, - *_LATE_INTERACTION_EMBEDDING_MODELS, + if any( + ( + FastEmbedMisc.is_supported_text_model(model_name), + FastEmbedMisc.is_supported_sparse_model(model_name), + FastEmbedMisc.is_supported_late_interaction_text_model(model_name), + ) ): return models.Document(model=model_name, text=value, options=options) - if model_name in _IMAGE_EMBEDDING_MODELS: + if FastEmbedMisc.is_supported_image_model(model_name): return models.Image(model=model_name, image=value, options=options) - if model_name in _LATE_INTERACTION_MULTIMODAL_EMBEDDING_MODELS: + if FastEmbedMisc.is_supported_late_interaction_multimodal_model(model_name): raise ValueError(f"{model_name} does not support `InferenceObject` interface") raise ValueError(f"{model_name} is not among supported models") diff --git a/qdrant_client/fastembed_common.py b/qdrant_client/fastembed_common.py index 4b42be1b..6744aed7 100644 --- a/qdrant_client/fastembed_common.py +++ b/qdrant_client/fastembed_common.py @@ -7,24 +7,268 @@ from qdrant_client.http import models try: from fastembed import ( - SparseTextEmbedding, TextEmbedding, - LateInteractionTextEmbedding, + SparseTextEmbedding, ImageEmbedding, + LateInteractionTextEmbedding, LateInteractionMultimodalEmbedding, ) - from fastembed.text.multitask_embedding import JinaEmbeddingV3 as _MultitaskTextEmbedding from fastembed.common import OnnxProvider, ImageInput except ImportError: TextEmbedding = None SparseTextEmbedding = None - OnnxProvider = None + ImageEmbedding = None LateInteractionTextEmbedding = None LateInteractionMultimodalEmbedding = None - ImageEmbedding = None + OnnxProvider = None ImageInput = None +class QueryResponse(BaseModel, extra="forbid"): # type: ignore + id: Union[str, int] + embedding: Optional[list[float]] + sparse_embedding: Optional[SparseVector] = Field(default=None) + metadata: dict[str, Any] + document: str + score: float + + +class FastEmbedMisc: + IS_INSTALLED: bool = False + _TEXT_MODELS: set[str] = set() + _IMAGE_MODELS: set[str] = set() + _LATE_INTERACTION_TEXT_MODELS: set[str] = set() + _LATE_INTERACTION_MULTIMODAL_MODELS: set[str] = set() + _SPARSE_MODELS: set[str] = set() + + @classmethod + def is_installed(cls) -> bool: + if cls.IS_INSTALLED: + return cls.IS_INSTALLED + + try: + from fastembed import ( + SparseTextEmbedding, + TextEmbedding, + ImageEmbedding, + LateInteractionMultimodalEmbedding, + LateInteractionTextEmbedding, + ) + + assert len(SparseTextEmbedding.list_supported_models()) > 0 + assert len(TextEmbedding.list_supported_models()) > 0 + assert len(ImageEmbedding.list_supported_models()) > 0 + assert len(LateInteractionTextEmbedding.list_supported_models()) > 0 + assert len(LateInteractionMultimodalEmbedding.list_supported_models()) > 0 + cls.IS_INSTALLED = True + except ImportError: + cls.IS_INSTALLED = False + + return cls.IS_INSTALLED + + @classmethod + def import_fastembed(cls) -> None: + if cls.IS_INSTALLED: + return + + # If it's not, ask the user to install it + raise ImportError( + "fastembed is not installed." + " Please install it to enable fast vector indexing with `pip install fastembed`." + ) + + @classmethod + def list_text_models(cls) -> dict[str, tuple[int, models.Distance]]: + """Lists the supported dense text models. + + Requires invocation of TextEmbedding.list_supported_models() to support custom models. + + Returns: + dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics. + """ + return ( + { + model["model"]: (model["dim"], models.Distance.COSINE) + for model in TextEmbedding.list_supported_models() + } + if TextEmbedding + else {} + ) + + @classmethod + def list_image_models(cls) -> dict[str, tuple[int, models.Distance]]: + """Lists the supported image dense models. + + Custom image models are not supported yet, but calls to ImageEmbedding.list_supported_models() is done each + time in order for preserving the same style as with TextEmbedding. + + Returns: + dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics. + """ + return ( + { + model["model"]: (model["dim"], models.Distance.COSINE) + for model in ImageEmbedding.list_supported_models() + } + if ImageEmbedding + else {} + ) + + @classmethod + def list_late_interaction_text_models(cls) -> dict[str, tuple[int, models.Distance]]: + """Lists the supported late interaction text models. + + Custom late interaction models are not supported yet, but calls to + LateInteractionTextEmbedding.list_supported_models() + is done each time in order for preserving the same style as with TextEmbedding. + + Returns: + dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics. + """ + return ( + { + model["model"]: (model["dim"], models.Distance.COSINE) + for model in LateInteractionTextEmbedding.list_supported_models() + } + if LateInteractionTextEmbedding + else {} + ) + + @classmethod + def list_late_interaction_multimodal_models(cls) -> dict[str, tuple[int, models.Distance]]: + """Lists the supported late interaction multimodal models. + + Custom late interaction multimodal models are not supported yet, but calls to + LateInteractionMultimodalEmbedding.list_supported_models() + is done each time in order for preserving the same style as with TextEmbedding. + + Returns: + dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics. + """ + return ( + { + model["model"]: (model["dim"], models.Distance.COSINE) + for model in LateInteractionMultimodalEmbedding.list_supported_models() + } + if LateInteractionMultimodalEmbedding + else {} + ) + + @classmethod + def list_sparse_models(cls) -> dict[str, dict[str, Any]]: + """Lists the supported sparse models. + + Custom sparse models are not supported yet, but calls to + SparseTextEmbedding.list_supported_models() + is done each time in order for preserving the same style as with TextEmbedding. + + Returns: + dict[str, dict[str, Any]]: A dict of model names and their descriptions. + """ + descriptions = {} + if SparseTextEmbedding: + for description in SparseTextEmbedding.list_supported_models(): + descriptions[description.pop("model")] = description + return descriptions + + @classmethod + def is_supported_text_model(cls, model_name: str) -> bool: + """Checks if the model is supported by fastembed. + + Args: + model_name (str): The name of the model to check. + + Returns: + bool: True if the model is supported, False otherwise. + """ + if model_name.lower() in cls._TEXT_MODELS: + return True + # update cached list in case custom models were added + cls._TEXT_MODELS = {model.lower() for model in cls.list_text_models()} + if model_name.lower() in cls._TEXT_MODELS: + return True + return False + + @classmethod + def is_supported_image_model(cls, model_name: str) -> bool: + """Checks if the model is supported by fastembed. + + Args: + model_name (str): The name of the model to check. + + Returns: + bool: True if the model is supported, False otherwise. + """ + if model_name.lower() in cls._IMAGE_MODELS: + return True + # update cached list in case custom models were added + cls._IMAGE_MODELS = {model.lower() for model in cls.list_image_models()} + if model_name.lower() in cls._IMAGE_MODELS: + return True + return False + + @classmethod + def is_supported_late_interaction_text_model(cls, model_name: str) -> bool: + """Checks if the model is supported by fastembed. + + Args: + model_name (str): The name of the model to check. + + Returns: + bool: True if the model is supported, False otherwise. + """ + if model_name.lower() in cls._LATE_INTERACTION_TEXT_MODELS: + return True + # update cached list in case custom models were added + cls._LATE_INTERACTION_TEXT_MODELS = { + model.lower() for model in cls.list_late_interaction_text_models() + } + if model_name.lower() in cls._LATE_INTERACTION_TEXT_MODELS: + return True + return False + + @classmethod + def is_supported_late_interaction_multimodal_model(cls, model_name: str) -> bool: + """Checks if the model is supported by fastembed. + + Args: + model_name (str): The name of the model to check. + + Returns: + bool: True if the model is supported, False otherwise. + """ + if model_name.lower() in cls._LATE_INTERACTION_MULTIMODAL_MODELS: + return True + # update cached list in case custom models were added + cls._LATE_INTERACTION_MULTIMODAL_MODELS = { + model.lower() for model in cls.list_late_interaction_multimodal_models() + } + if model_name.lower() in cls._LATE_INTERACTION_MULTIMODAL_MODELS: + return True + return False + + @classmethod + def is_supported_sparse_model(cls, model_name: str) -> bool: + """Checks if the model is supported by fastembed. + + Args: + model_name (str): The name of the model to check. + + Returns: + bool: True if the model is supported, False otherwise. + """ + if model_name.lower() in cls._SPARSE_MODELS: + return True + # update cached list in case custom models were added + cls._SPARSE_MODELS = {model.lower() for model in cls.list_sparse_models()} + if model_name.lower() in cls._SPARSE_MODELS: + return True + return False + + +# region deprecated +# prefer using methods builtin into QdrantClient, e.g. list_supported_text_models, list_supported_idf_models, etc. + SUPPORTED_EMBEDDING_MODELS: dict[str, tuple[int, models.Distance]] = ( { model["model"]: (model["dim"], models.Distance.COSINE) @@ -76,12 +320,4 @@ _LATE_INTERACTION_MULTIMODAL_EMBEDDING_MODELS: dict[str, tuple[int, models.Dista if LateInteractionMultimodalEmbedding else {} ) - - -class QueryResponse(BaseModel, extra="forbid"): # type: ignore - id: Union[str, int] - embedding: Optional[list[float]] - sparse_embedding: Optional[SparseVector] = Field(default=None) - metadata: dict[str, Any] - document: str - score: float +# endregion diff --git a/qdrant_client/qdrant_fastembed.py b/qdrant_client/qdrant_fastembed.py index 727c383c..b12312a5 100644 --- a/qdrant_client/qdrant_fastembed.py +++ b/qdrant_client/qdrant_fastembed.py @@ -16,19 +16,16 @@ 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, - LateInteractionTextEmbedding, - ImageEmbedding, SparseTextEmbedding, - SUPPORTED_EMBEDDING_MODELS, - SUPPORTED_SPARSE_EMBEDDING_MODELS, - _LATE_INTERACTION_EMBEDDING_MODELS, - _IMAGE_EMBEDDING_MODELS, IDF_EMBEDDING_MODELS, - OnnxProvider, ) +# endregion class QdrantFastembedMixin(QdrantBase): @@ -41,18 +38,54 @@ class QdrantFastembedMixin(QdrantBase): self._sparse_embedding_model_name: Optional[str] = None self._model_embedder = ModelEmbedder(parser=parser, **kwargs) - try: - from fastembed import SparseTextEmbedding, TextEmbedding - - assert len(SparseTextEmbedding.list_supported_models()) > 0 - assert len(TextEmbedding.list_supported_models()) > 0 - - self.__class__._FASTEMBED_INSTALLED = True - except ImportError: - self.__class__._FASTEMBED_INSTALLED = False - + self.__class__._FASTEMBED_INSTALLED = FastEmbedMisc.is_installed() super().__init__(**kwargs) + @classmethod + def list_text_models(cls) -> dict[str, tuple[int, models.Distance]]: + """Lists the supported dense text models. + + Returns: + dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics. + """ + return FastEmbedMisc.list_text_models() + + @classmethod + def list_image_models(cls) -> dict[str, tuple[int, models.Distance]]: + """Lists the supported image dense models. + + Returns: + dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics. + """ + return FastEmbedMisc.list_image_models() + + @classmethod + def list_late_interaction_text_models(cls) -> dict[str, tuple[int, models.Distance]]: + """Lists the supported late interaction text models. + + Returns: + dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics. + """ + return FastEmbedMisc.list_late_interaction_text_models() + + @classmethod + def list_late_interaction_multimodal_models(cls) -> dict[str, tuple[int, models.Distance]]: + """Lists the supported late interaction multimodal models. + + Returns: + dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics. + """ + return FastEmbedMisc.list_late_interaction_multimodal_models() + + @classmethod + def list_sparse_models(cls) -> dict[str, dict[str, Any]]: + """Lists the supported sparse text models. + + Returns: + dict[str, dict[str, Any]]: A dict of model names and their descriptions. + """ + return FastEmbedMisc.list_sparse_models() + @property def embedding_model_name(self) -> str: if self._embedding_model_name is None: @@ -174,31 +207,20 @@ class QdrantFastembedMixin(QdrantBase): ) self._sparse_embedding_model_name = embedding_model_name - @classmethod - def _import_fastembed(cls) -> None: - if cls._FASTEMBED_INSTALLED: - return - - # If it's not, ask the user to install it - raise ImportError( - "fastembed is not installed." - " Please install it to enable fast vector indexing with `pip install fastembed`." - ) - @classmethod def _get_model_params(cls, model_name: str) -> tuple[int, models.Distance]: - cls._import_fastembed() + FastEmbedMisc.import_fastembed() - if model_name in SUPPORTED_EMBEDDING_MODELS: - return SUPPORTED_EMBEDDING_MODELS[model_name] + for descriptions in ( + FastEmbedMisc.list_text_models(), + FastEmbedMisc.list_image_models(), + FastEmbedMisc.list_late_interaction_text_models(), + FastEmbedMisc.list_late_interaction_multimodal_models(), + ): + if params := descriptions.get(model_name): + return params - if model_name in _LATE_INTERACTION_EMBEDDING_MODELS: - return _LATE_INTERACTION_EMBEDDING_MODELS[model_name] - - if model_name in _IMAGE_EMBEDDING_MODELS: - return _IMAGE_EMBEDDING_MODELS[model_name] - - if model_name in SUPPORTED_SPARSE_EMBEDDING_MODELS: + if model_name in FastEmbedMisc.list_sparse_models(): raise ValueError( "Sparse embeddings do not return fixed embedding size and distance type" ) @@ -214,7 +236,7 @@ class QdrantFastembedMixin(QdrantBase): deprecated: bool = False, **kwargs: Any, ) -> "TextEmbedding": - self._import_fastembed() + FastEmbedMisc.import_fastembed() return self._model_embedder.embedder.get_or_init_model( model_name=model_name, @@ -234,7 +256,7 @@ class QdrantFastembedMixin(QdrantBase): deprecated: bool = False, **kwargs: Any, ) -> "SparseTextEmbedding": - self._import_fastembed() + FastEmbedMisc.import_fastembed() return self._model_embedder.embedder.get_or_init_sparse_model( model_name=model_name, @@ -245,41 +267,6 @@ class QdrantFastembedMixin(QdrantBase): **kwargs, ) - def _get_or_init_late_interaction_model( - self, - model_name: str, - cache_dir: Optional[str] = None, - threads: Optional[int] = None, - providers: Optional[Sequence["OnnxProvider"]] = None, - **kwargs: Any, - ) -> "LateInteractionTextEmbedding": - self._import_fastembed() - return self._model_embedder.embedder.get_or_init_late_interaction_model( - model_name=model_name, - cache_dir=cache_dir, - threads=threads, - providers=providers, - **kwargs, - ) - - def _get_or_init_image_model( - self, - model_name: str, - cache_dir: Optional[str] = None, - threads: Optional[int] = None, - providers: Optional[Sequence["OnnxProvider"]] = None, - **kwargs: Any, - ) -> "ImageEmbedding": - self._import_fastembed() - - return self._model_embedder.embedder.get_or_init_image_model( - model_name=model_name, - cache_dir=cache_dir, - threads=threads, - providers=providers, - **kwargs, - ) - def _embed_documents( self, documents: Iterable[str], @@ -453,20 +440,19 @@ class QdrantFastembedMixin(QdrantBase): self, model_name: Optional[str] = None, ) -> int: - """ - Get the size of the embeddings produced by the specified model. + """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: optional, the name of the model to get the embedding size for. If None, the default model will + be used. 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 - if model_name in SUPPORTED_SPARSE_EMBEDDING_MODELS: - raise ValueError( - f"Sparse embeddings do not have a fixed embedding size. Current model: {model_name}" - ) embeddings_size, _ = self._get_model_params(model_name=model_name) return embeddings_size @@ -884,6 +870,8 @@ class QdrantFastembedMixin(QdrantBase): is_query: bool = False, batch_size: Optional[int] = None, ) -> Iterable[BaseModel]: + FastEmbedMisc.import_fastembed() + yield from self._model_embedder.embed_models( raw_models=raw_models, is_query=is_query, @@ -896,6 +884,7 @@ class QdrantFastembedMixin(QdrantBase): batch_size: Optional[int] = None, parallel: Optional[int] = None, ) -> Iterable[BaseModel]: + FastEmbedMisc.import_fastembed() yield from self._model_embedder.embed_models_strict( raw_models=raw_models, batch_size=batch_size or self.DEFAULT_BATCH_SIZE, diff --git a/tests/test_fastembed.py b/tests/test_fastembed.py index d089dd9b..29017ef0 100644 --- a/tests/test_fastembed.py +++ b/tests/test_fastembed.py @@ -217,5 +217,7 @@ def test_get_embedding_size(): assert local_client.get_embedding_size(model_name="colbert-ir/colbertv2.0") == 128 - with pytest.raises(ValueError, match="Sparse embeddings do not have a fixed embedding size."): + with pytest.raises( + ValueError, match="Sparse embeddings do not return fixed embedding size and distance type" + ): local_client.get_embedding_size(model_name="Qdrant/bm25")