mirror of
https://github.com/qdrant/qdrant-client.git
synced 2026-07-23 11:11:01 -05:00
wip: builtin bm25 support (#1060)
* wip: builtin bm25 support * remove debug print * new: add builtin embedder, add server version check for inference * fix: regen async * fix: fix fastembed import * fix: regen async * tests: add tests for builtin inference * fix: fix fastembed checks, improve exceptions
This commit is contained in:
@@ -105,13 +105,12 @@ class AsyncQdrantClient(AsyncQdrantFastembedMixin):
|
||||
if key not in ("self", "__class__", "kwargs")
|
||||
}
|
||||
self._init_options.update({k: v for (k, v) in kwargs.items()})
|
||||
self._inference_inspector = Inspector()
|
||||
super().__init__(parser=self._inference_inspector.parser, **kwargs)
|
||||
self._client: AsyncQdrantBase
|
||||
if sum([param is not None for param in (location, url, host, path)]) > 1:
|
||||
raise ValueError(
|
||||
"Only one of <location>, <url>, <host> or <path> should be specified."
|
||||
)
|
||||
self._client: AsyncQdrantBase
|
||||
server_version = None
|
||||
if location == ":memory:":
|
||||
self._client = AsyncQdrantLocal(
|
||||
location=location, force_disable_check_same_thread=force_disable_check_same_thread
|
||||
@@ -138,12 +137,19 @@ class AsyncQdrantClient(AsyncQdrantFastembedMixin):
|
||||
check_compatibility=check_compatibility,
|
||||
**kwargs,
|
||||
)
|
||||
server_version = self._client.server_version
|
||||
if isinstance(self._client, AsyncQdrantLocal) and cloud_inference:
|
||||
raise ValueError(
|
||||
"Cloud inference is not supported for local Qdrant, consider using FastEmbed or switch to Qdrant Cloud"
|
||||
)
|
||||
self.cloud_inference = cloud_inference
|
||||
self.local_inference_batch_size = local_inference_batch_size
|
||||
self._inference_inspector = Inspector()
|
||||
super().__init__(
|
||||
parser=self._inference_inspector.parser,
|
||||
is_local_mode=isinstance(self._client, AsyncQdrantLocal),
|
||||
server_version=server_version,
|
||||
)
|
||||
|
||||
async def close(self, grpc_grace: Optional[float] = None, **kwargs: Any) -> None:
|
||||
"""Closes the connection to Qdrant
|
||||
|
||||
@@ -18,6 +18,7 @@ from pydantic import BaseModel
|
||||
from qdrant_client import grpc
|
||||
from qdrant_client.common.client_warnings import show_warning
|
||||
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
|
||||
@@ -39,12 +40,16 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase):
|
||||
DEFAULT_BATCH_SIZE = 8
|
||||
_FASTEMBED_INSTALLED: bool
|
||||
|
||||
def __init__(self, parser: ModelSchemaParser, **kwargs: Any):
|
||||
def __init__(
|
||||
self, parser: ModelSchemaParser, is_local_mode: bool, server_version: Optional[str]
|
||||
):
|
||||
self.__class__._FASTEMBED_INSTALLED = FastEmbedMisc.is_installed()
|
||||
self._embedding_model_name: Optional[str] = None
|
||||
self._sparse_embedding_model_name: Optional[str] = None
|
||||
self._model_embedder = ModelEmbedder(parser=parser, **kwargs)
|
||||
self.__class__._FASTEMBED_INSTALLED = FastEmbedMisc.is_installed()
|
||||
super().__init__(**kwargs)
|
||||
self._model_embedder = ModelEmbedder(
|
||||
parser=parser, is_local_mode=is_local_mode, server_version=server_version
|
||||
)
|
||||
super().__init__()
|
||||
|
||||
@classmethod
|
||||
async def list_text_models(cls) -> dict[str, tuple[int, models.Distance]]:
|
||||
@@ -238,6 +243,7 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase):
|
||||
**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,
|
||||
@@ -257,6 +263,7 @@ class AsyncQdrantFastembedMixin(AsyncQdrantBase):
|
||||
**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,
|
||||
@@ -800,7 +807,6 @@ 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,
|
||||
@@ -813,7 +819,6 @@ 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,
|
||||
|
||||
@@ -171,21 +171,22 @@ class AsyncQdrantRemote(AsyncQdrantBase):
|
||||
self._grpc_snapshots_client: Optional[grpc.SnapshotsStub] = None
|
||||
self._grpc_root_client: Optional[grpc.QdrantStub] = None
|
||||
self._closed: bool = False
|
||||
self.server_version = None
|
||||
if check_compatibility:
|
||||
try:
|
||||
client_version = importlib.metadata.version("qdrant-client")
|
||||
server_version = get_server_version(
|
||||
self.server_version = get_server_version(
|
||||
self.rest_uri, self._rest_headers, self._rest_args.get("auth")
|
||||
)
|
||||
if not server_version:
|
||||
if not self.server_version:
|
||||
show_warning(
|
||||
message="Failed to obtain server version. Unable to check client-server compatibility. Set check_compatibility=False to skip version check.",
|
||||
category=UserWarning,
|
||||
stacklevel=4,
|
||||
)
|
||||
elif not is_compatible(client_version, server_version):
|
||||
elif not is_compatible(client_version, self.server_version):
|
||||
show_warning(
|
||||
message=f"Qdrant client version {client_version} is incompatible with server version {server_version}. Major versions should match and minor version difference must not exceed 1. Set check_compatibility=False to skip version check.",
|
||||
message=f"Qdrant client version {client_version} is incompatible with server version {self.server_version}. Major versions should match and minor version difference must not exceed 1. Set check_compatibility=False to skip version check.",
|
||||
category=UserWarning,
|
||||
stacklevel=4,
|
||||
)
|
||||
|
||||
94
qdrant_client/embed/builtin_embedder.py
Normal file
94
qdrant_client/embed/builtin_embedder.py
Normal file
@@ -0,0 +1,94 @@
|
||||
from typing import Optional, Any
|
||||
|
||||
from qdrant_client.http import models
|
||||
from qdrant_client.embed.models import NumericVector
|
||||
|
||||
|
||||
class BuiltinEmbedder:
|
||||
_SUPPORTED_MODELS = ("Qdrant/Bm25",)
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
def embed(
|
||||
self,
|
||||
model_name: str,
|
||||
texts: Optional[list[str]] = None,
|
||||
options: Optional[dict[str, Any]] = None,
|
||||
**kwargs: Any,
|
||||
) -> NumericVector:
|
||||
if texts is None:
|
||||
if "images" in kwargs:
|
||||
raise ValueError(
|
||||
"Image processing is only available with cloud inference of FastEmbed"
|
||||
)
|
||||
|
||||
raise ValueError("Texts must be provided for the inference")
|
||||
|
||||
if not self.is_supported_sparse_model(model_name):
|
||||
raise ValueError(
|
||||
f"Model {model_name} is not supported in {self.__class__.__name__}. "
|
||||
f"Did you forget to enable cloud inference or install FastEmbed for local inference?"
|
||||
)
|
||||
|
||||
return [models.Document(text=text, options=options, model=model_name) for text in texts]
|
||||
|
||||
@classmethod
|
||||
def is_supported_text_model(cls, model_name: str) -> bool:
|
||||
"""Mock embedder interface, only sparse text model Qdrant/Bm25 is supported
|
||||
|
||||
Args:
|
||||
model_name (str): The name of the model to check.
|
||||
|
||||
Returns:
|
||||
bool: True if the model is supported, False otherwise.
|
||||
"""
|
||||
return False # currently only Qdrant/Bm25 is supported
|
||||
|
||||
@classmethod
|
||||
def is_supported_image_model(cls, model_name: str) -> bool:
|
||||
"""Mock embedder interface, only sparse text model Qdrant/Bm25 is supported
|
||||
|
||||
Args:
|
||||
model_name (str): The name of the model to check.
|
||||
|
||||
Returns:
|
||||
bool: True if the model is supported, False otherwise.
|
||||
"""
|
||||
return False # currently only Qdrant/Bm25 is supported
|
||||
|
||||
@classmethod
|
||||
def is_supported_late_interaction_text_model(cls, model_name: str) -> bool:
|
||||
"""Mock embedder interface, only sparse text model Qdrant/Bm25 is supported
|
||||
|
||||
Args:
|
||||
model_name (str): The name of the model to check.
|
||||
|
||||
Returns:
|
||||
bool: True if the model is supported, False otherwise.
|
||||
"""
|
||||
return False # currently only Qdrant/Bm25 is supported
|
||||
|
||||
@classmethod
|
||||
def is_supported_late_interaction_multimodal_model(cls, model_name: str) -> bool:
|
||||
"""Mock embedder interface, only sparse text model Qdrant/Bm25 is supported
|
||||
|
||||
Args:
|
||||
model_name (str): The name of the model to check.
|
||||
|
||||
Returns:
|
||||
bool: True if the model is supported, False otherwise.
|
||||
"""
|
||||
return False # currently only Qdrant/Bm25 is supported
|
||||
|
||||
@classmethod
|
||||
def is_supported_sparse_model(cls, model_name: str) -> bool:
|
||||
"""Checks if the model is supported. Only `Qdrant/Bm25` is supported
|
||||
|
||||
Args:
|
||||
model_name (str): The name of the model to check.
|
||||
|
||||
Returns:
|
||||
bool: True if the model is supported, False otherwise.
|
||||
"""
|
||||
return model_name.lower() in [model.lower() for model in cls._SUPPORTED_MODELS]
|
||||
@@ -385,3 +385,63 @@ class Embedder:
|
||||
for embedding in embedding_model_inst.embed(images=images, batch_size=batch_size)
|
||||
]
|
||||
return embeddings
|
||||
|
||||
@classmethod
|
||||
def is_supported_text_model(cls, model_name: str) -> bool:
|
||||
"""Check if 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.
|
||||
"""
|
||||
return FastEmbedMisc.is_supported_text_model(model_name)
|
||||
|
||||
@classmethod
|
||||
def is_supported_image_model(cls, model_name: str) -> bool:
|
||||
"""Check if 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.
|
||||
"""
|
||||
return FastEmbedMisc.is_supported_image_model(model_name)
|
||||
|
||||
@classmethod
|
||||
def is_supported_late_interaction_text_model(cls, model_name: str) -> bool:
|
||||
"""Check if 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.
|
||||
"""
|
||||
return FastEmbedMisc.is_supported_late_interaction_text_model(model_name)
|
||||
|
||||
@classmethod
|
||||
def is_supported_late_interaction_multimodal_model(cls, model_name: str) -> bool:
|
||||
"""Check if 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.
|
||||
"""
|
||||
return FastEmbedMisc.is_supported_late_interaction_multimodal_model(model_name)
|
||||
|
||||
@classmethod
|
||||
def is_supported_sparse_model(cls, model_name: str) -> bool:
|
||||
"""Check if 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.
|
||||
"""
|
||||
return FastEmbedMisc.is_supported_sparse_model(model_name)
|
||||
|
||||
@@ -6,6 +6,7 @@ from typing import Optional, Union, Iterable, Any, Type, get_args
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from qdrant_client.embed.builtin_embedder import BuiltinEmbedder
|
||||
from qdrant_client.http import models
|
||||
from qdrant_client.embed.common import INFERENCE_OBJECT_TYPES
|
||||
from qdrant_client.embed.embed_inspector import InspectorEmbed
|
||||
@@ -42,11 +43,46 @@ class ModelEmbedderWorker(Worker):
|
||||
class ModelEmbedder:
|
||||
MAX_INTERNAL_BATCH_SIZE = 64
|
||||
|
||||
def __init__(self, parser: Optional[ModelSchemaParser] = None, **kwargs: Any):
|
||||
def __init__(
|
||||
self,
|
||||
parser: Optional[ModelSchemaParser] = None,
|
||||
is_local_mode: bool = False,
|
||||
server_version: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
self._batch_accumulator: dict[str, list[INFERENCE_OBJECT_TYPES]] = {}
|
||||
self._embed_storage: dict[str, list[NumericVector]] = {}
|
||||
self._embed_inspector = InspectorEmbed(parser=parser)
|
||||
self.embedder = Embedder(**kwargs)
|
||||
self._is_builtin_embedder_available = self._check_builtin_embedder_availability(
|
||||
is_local_mode, server_version
|
||||
)
|
||||
self.embedder = (
|
||||
Embedder(**kwargs) if FastEmbedMisc.is_installed() else BuiltinEmbedder(**kwargs)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _check_builtin_embedder_availability(
|
||||
is_local_mode: bool, server_version: Optional[str]
|
||||
) -> bool:
|
||||
if is_local_mode:
|
||||
return False
|
||||
|
||||
if (
|
||||
server_version is None
|
||||
): # failed to detect server version, it might happen due to security or network
|
||||
# problems even on supported server versions, so we are not blocking usage of BuiltinEmbedder.
|
||||
return True
|
||||
|
||||
try:
|
||||
major, minor, patch = server_version.split(".")
|
||||
patch = patch.split("-")[0]
|
||||
|
||||
if (int(major), int(minor), int(patch)) >= (1, 15, 3):
|
||||
return True
|
||||
|
||||
return False
|
||||
except Exception:
|
||||
return True
|
||||
|
||||
def embed_models(
|
||||
self,
|
||||
@@ -65,6 +101,9 @@ class ModelEmbedder:
|
||||
Returns:
|
||||
list[BaseModel]: models with embedded fields
|
||||
"""
|
||||
if not self._is_builtin_embedder_available:
|
||||
FastEmbedMisc.import_fastembed() # fail fast if fastembed is required
|
||||
|
||||
if isinstance(raw_models, BaseModel):
|
||||
raw_models = [raw_models]
|
||||
for raw_models_batch in iter_batch(raw_models, batch_size):
|
||||
@@ -91,13 +130,21 @@ class ModelEmbedder:
|
||||
Returns:
|
||||
Iterable[Union[dict[str, BaseModel], BaseModel]]: models with embedded fields
|
||||
"""
|
||||
if not self._is_builtin_embedder_available:
|
||||
FastEmbedMisc.import_fastembed() # fail fast if fastembed is required
|
||||
|
||||
is_small = False
|
||||
|
||||
if isinstance(raw_models, list):
|
||||
if len(raw_models) < batch_size:
|
||||
is_small = True
|
||||
|
||||
if parallel is None or parallel == 1 or is_small:
|
||||
if (
|
||||
isinstance(self.embedder, BuiltinEmbedder)
|
||||
or parallel is None
|
||||
or parallel == 1
|
||||
or is_small
|
||||
):
|
||||
for batch in iter_batch(raw_models, batch_size):
|
||||
yield from self.embed_models_batch(batch, inference_batch_size=batch_size)
|
||||
else:
|
||||
@@ -138,6 +185,9 @@ class ModelEmbedder:
|
||||
Returns:
|
||||
Iterable[BaseModel]: models with embedded fields
|
||||
"""
|
||||
if not self._is_builtin_embedder_available:
|
||||
FastEmbedMisc.import_fastembed() # fail fast if fastembed is required
|
||||
|
||||
for raw_model in raw_models:
|
||||
self._process_model(raw_model, is_query=is_query, accumulating=True)
|
||||
|
||||
@@ -379,14 +429,20 @@ class ModelEmbedder:
|
||||
for model in self._batch_accumulator:
|
||||
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),
|
||||
self.embedder.is_supported_text_model(model),
|
||||
self.embedder.is_supported_sparse_model(model),
|
||||
self.embedder.is_supported_late_interaction_text_model(model),
|
||||
self.embedder.is_supported_image_model(model),
|
||||
self.embedder.is_supported_late_interaction_multimodal_model(model),
|
||||
)
|
||||
):
|
||||
raise ValueError(f"{model} is not among supported models")
|
||||
if isinstance(self.embedder, BuiltinEmbedder):
|
||||
raise ValueError(
|
||||
f"{model} is not among supported models. "
|
||||
f"Have you forgotten to set `cloud_inference` or install `fastembed` for local inference?"
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"{model} is not among supported models")
|
||||
|
||||
for model, data in self._batch_accumulator.items():
|
||||
self._embed_storage[model] = embed(
|
||||
@@ -405,8 +461,7 @@ class ModelEmbedder:
|
||||
"""
|
||||
return self._embed_storage[model_name].pop(0)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_inference_object(data: models.VectorStruct) -> models.VectorStruct:
|
||||
def _resolve_inference_object(self, data: models.VectorStruct) -> models.VectorStruct:
|
||||
"""Resolve inference object into a model
|
||||
|
||||
Args:
|
||||
@@ -425,15 +480,15 @@ class ModelEmbedder:
|
||||
options = data.options
|
||||
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),
|
||||
self.embedder.is_supported_text_model(model_name),
|
||||
self.embedder.is_supported_sparse_model(model_name),
|
||||
self.embedder.is_supported_late_interaction_text_model(model_name),
|
||||
)
|
||||
):
|
||||
return models.Document(model=model_name, text=value, options=options)
|
||||
if FastEmbedMisc.is_supported_image_model(model_name):
|
||||
if self.embedder.is_supported_image_model(model_name):
|
||||
return models.Image(model=model_name, image=value, options=options)
|
||||
if FastEmbedMisc.is_supported_late_interaction_multimodal_model(model_name):
|
||||
if self.embedder.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")
|
||||
|
||||
@@ -109,20 +109,13 @@ class QdrantClient(QdrantFastembedMixin):
|
||||
}
|
||||
self._init_options.update({k: v for k, v in kwargs.items()})
|
||||
|
||||
self._inference_inspector = Inspector()
|
||||
super().__init__(
|
||||
parser=self._inference_inspector.parser, **kwargs
|
||||
) # If we want to pass any kwargs to the parent class or ignore unexpected kwargs,
|
||||
# we will need to pop them from **kwargs. Otherwise, they might be passed to QdrantRemote as httpx kwargs.
|
||||
# Httpx has specific set of params, which it accepts and will raise an error if it receives any other params.
|
||||
|
||||
self._client: QdrantBase
|
||||
|
||||
if sum([param is not None for param in (location, url, host, path)]) > 1:
|
||||
raise ValueError(
|
||||
"Only one of <location>, <url>, <host> or <path> should be specified."
|
||||
)
|
||||
self._client: QdrantBase
|
||||
|
||||
server_version = None
|
||||
if location == ":memory:":
|
||||
self._client = QdrantLocal(
|
||||
location=location,
|
||||
@@ -151,6 +144,7 @@ class QdrantClient(QdrantFastembedMixin):
|
||||
check_compatibility=check_compatibility,
|
||||
**kwargs,
|
||||
)
|
||||
server_version = self._client.server_version
|
||||
|
||||
if isinstance(self._client, QdrantLocal) and cloud_inference:
|
||||
raise ValueError(
|
||||
@@ -159,6 +153,16 @@ class QdrantClient(QdrantFastembedMixin):
|
||||
self.cloud_inference = cloud_inference
|
||||
self.local_inference_batch_size = local_inference_batch_size
|
||||
|
||||
self._inference_inspector = Inspector()
|
||||
super().__init__(
|
||||
parser=self._inference_inspector.parser,
|
||||
is_local_mode=isinstance(self._client, QdrantLocal),
|
||||
server_version=server_version,
|
||||
) # If we'd like to pass any kwargs to the parent class or ignore unexpected kwargs,
|
||||
# we will need to pop them from **kwargs and call super().__init__ before creating QdrantRemote instance.
|
||||
# Otherwise, they might be passed to QdrantRemote as httpx kwargs.
|
||||
# Httpx has specific set of params, which it accepts and will raise an error if it receives any other params.
|
||||
|
||||
def __del__(self) -> None:
|
||||
self.close()
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ from pydantic import BaseModel
|
||||
from qdrant_client import grpc
|
||||
from qdrant_client.common.client_warnings import show_warning
|
||||
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
|
||||
@@ -33,13 +34,17 @@ class QdrantFastembedMixin(QdrantBase):
|
||||
DEFAULT_BATCH_SIZE = 8
|
||||
_FASTEMBED_INSTALLED: bool
|
||||
|
||||
def __init__(self, parser: ModelSchemaParser, **kwargs: Any):
|
||||
def __init__(
|
||||
self, parser: ModelSchemaParser, is_local_mode: bool, server_version: Optional[str]
|
||||
):
|
||||
self.__class__._FASTEMBED_INSTALLED = FastEmbedMisc.is_installed()
|
||||
self._embedding_model_name: Optional[str] = None
|
||||
self._sparse_embedding_model_name: Optional[str] = None
|
||||
self._model_embedder = ModelEmbedder(parser=parser, **kwargs)
|
||||
|
||||
self.__class__._FASTEMBED_INSTALLED = FastEmbedMisc.is_installed()
|
||||
super().__init__(**kwargs)
|
||||
self._model_embedder = ModelEmbedder(
|
||||
parser=parser, is_local_mode=is_local_mode, server_version=server_version
|
||||
)
|
||||
super().__init__()
|
||||
|
||||
@classmethod
|
||||
def list_text_models(cls) -> dict[str, tuple[int, models.Distance]]:
|
||||
@@ -238,6 +243,7 @@ class QdrantFastembedMixin(QdrantBase):
|
||||
) -> "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,
|
||||
@@ -257,7 +263,7 @@ class QdrantFastembedMixin(QdrantBase):
|
||||
**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,
|
||||
@@ -870,8 +876,6 @@ 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,
|
||||
@@ -884,7 +888,6 @@ 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,
|
||||
|
||||
@@ -206,24 +206,25 @@ class QdrantRemote(QdrantBase):
|
||||
|
||||
self._closed: bool = False
|
||||
|
||||
self.server_version = None
|
||||
if check_compatibility:
|
||||
try:
|
||||
client_version = importlib.metadata.version("qdrant-client")
|
||||
server_version = get_server_version(
|
||||
self.server_version = get_server_version(
|
||||
self.rest_uri, self._rest_headers, self._rest_args.get("auth")
|
||||
)
|
||||
|
||||
if not server_version:
|
||||
if not self.server_version:
|
||||
show_warning(
|
||||
message="Failed to obtain server version. Unable to check client-server compatibility."
|
||||
" Set check_compatibility=False to skip version check.",
|
||||
category=UserWarning,
|
||||
stacklevel=4,
|
||||
)
|
||||
elif not is_compatible(client_version, server_version):
|
||||
elif not is_compatible(client_version, self.server_version):
|
||||
show_warning(
|
||||
message=f"Qdrant client version {client_version} is incompatible with server "
|
||||
f"version {server_version}. Major versions should match and minor version difference "
|
||||
f"version {self.server_version}. Major versions should match and minor version difference "
|
||||
"must not exceed 1. Set check_compatibility=False to skip version check.",
|
||||
category=UserWarning,
|
||||
stacklevel=4,
|
||||
|
||||
178
tests/embed_tests/test_builtin_inference.py
Normal file
178
tests/embed_tests/test_builtin_inference.py
Normal file
@@ -0,0 +1,178 @@
|
||||
from typing import Any, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from qdrant_client import QdrantClient, models
|
||||
from qdrant_client.fastembed_common import FastEmbedMisc
|
||||
|
||||
COLLECTION_NAME = "inference_collection"
|
||||
MODEL_NAME = "Qdrant/Bm25"
|
||||
DEFAULT_VECTOR_NAME = "bm25"
|
||||
|
||||
|
||||
def prepare_collection(
|
||||
client: QdrantClient,
|
||||
collection_name: str,
|
||||
vectors_config: Optional[dict[str, Any]] = None,
|
||||
sparse_vectors_config: Optional[dict[str, Any]] = None,
|
||||
) -> None:
|
||||
if client.collection_exists(collection_name):
|
||||
client.delete_collection(collection_name)
|
||||
config = (
|
||||
{DEFAULT_VECTOR_NAME: models.SparseVectorParams(modifier=models.Modifier.IDF)}
|
||||
if sparse_vectors_config is None
|
||||
else sparse_vectors_config
|
||||
)
|
||||
client.create_collection(
|
||||
collection_name, vectors_config=vectors_config or {}, sparse_vectors_config=config
|
||||
)
|
||||
|
||||
|
||||
def test_bm25_inference():
|
||||
remote_client = QdrantClient()
|
||||
prepare_collection(remote_client, COLLECTION_NAME)
|
||||
local_client = QdrantClient(":memory:")
|
||||
prepare_collection(local_client, COLLECTION_NAME)
|
||||
|
||||
remote_client.upsert(
|
||||
COLLECTION_NAME,
|
||||
points=[
|
||||
models.PointStruct(
|
||||
id=1,
|
||||
vector={DEFAULT_VECTOR_NAME: models.Document(text="good text", model=MODEL_NAME)},
|
||||
)
|
||||
],
|
||||
)
|
||||
assert remote_client.count(collection_name=COLLECTION_NAME, exact=True).count == 1
|
||||
|
||||
# not calling is_installed() on purpose, since it changes `IS_INSTALLED` and might conceal a bug
|
||||
if FastEmbedMisc.IS_INSTALLED:
|
||||
local_client.upsert(
|
||||
COLLECTION_NAME,
|
||||
points=[
|
||||
models.PointStruct(
|
||||
id=1,
|
||||
vector={
|
||||
DEFAULT_VECTOR_NAME: models.Document(text="good text", model=MODEL_NAME)
|
||||
},
|
||||
)
|
||||
],
|
||||
)
|
||||
assert local_client.count(collection_name=COLLECTION_NAME, exact=True).count == 1
|
||||
else:
|
||||
# inference is done via builtin Qdrant bm25 in remote client, and is not available in local mode
|
||||
with pytest.raises(ImportError):
|
||||
local_client.upsert(
|
||||
COLLECTION_NAME,
|
||||
points=[
|
||||
models.PointStruct(
|
||||
id=1,
|
||||
vector={
|
||||
DEFAULT_VECTOR_NAME: models.Document(text="bad text", model=MODEL_NAME)
|
||||
},
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_bm25_inference_server_version(monkeypatch):
|
||||
server_version = "1.11.0"
|
||||
|
||||
def patched_get_server_version(*args, **kwargs):
|
||||
return server_version
|
||||
|
||||
monkeypatch.setattr(
|
||||
"qdrant_client.qdrant_remote.get_server_version", patched_get_server_version
|
||||
)
|
||||
remote_client = QdrantClient()
|
||||
prepare_collection(remote_client, COLLECTION_NAME)
|
||||
|
||||
if FastEmbedMisc.IS_INSTALLED:
|
||||
# inference is done via fastembed in both remote and local client
|
||||
remote_client.upsert(
|
||||
COLLECTION_NAME,
|
||||
points=[
|
||||
models.PointStruct(
|
||||
id=1,
|
||||
vector={
|
||||
DEFAULT_VECTOR_NAME: models.Document(text="good text", model=MODEL_NAME)
|
||||
},
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
assert remote_client.count(collection_name=COLLECTION_NAME, exact=True).count == 1
|
||||
else:
|
||||
with pytest.raises(ImportError):
|
||||
remote_client.upsert(
|
||||
COLLECTION_NAME,
|
||||
points=[
|
||||
models.PointStruct(
|
||||
id=1,
|
||||
vector={
|
||||
DEFAULT_VECTOR_NAME: models.Document(text="bad text", model=MODEL_NAME)
|
||||
},
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
server_version = None
|
||||
monkeypatch.setattr("qdrant_client.qdrant_remote.get_server_version", lambda: server_version)
|
||||
remote_client = QdrantClient()
|
||||
prepare_collection(remote_client, COLLECTION_NAME)
|
||||
|
||||
remote_client.upsert(
|
||||
COLLECTION_NAME,
|
||||
points=[
|
||||
models.PointStruct(
|
||||
id=1,
|
||||
vector={DEFAULT_VECTOR_NAME: models.Document(text="good text", model=MODEL_NAME)},
|
||||
)
|
||||
],
|
||||
)
|
||||
assert remote_client.count(collection_name=COLLECTION_NAME, exact=True).count == 1
|
||||
|
||||
server_version = "1.15.3"
|
||||
monkeypatch.setattr("qdrant_client.qdrant_remote.get_server_version", lambda: server_version)
|
||||
remote_client = QdrantClient()
|
||||
prepare_collection(remote_client, COLLECTION_NAME)
|
||||
remote_client.upsert(
|
||||
COLLECTION_NAME,
|
||||
points=[
|
||||
models.PointStruct(
|
||||
id=1,
|
||||
vector={DEFAULT_VECTOR_NAME: models.Document(text="good text", model=MODEL_NAME)},
|
||||
)
|
||||
],
|
||||
)
|
||||
assert remote_client.count(collection_name=COLLECTION_NAME, exact=True).count == 1
|
||||
|
||||
|
||||
def test_not_supported_builtin_inference_models():
|
||||
if FastEmbedMisc.is_installed():
|
||||
pytest.skip(reason="testing builtin inference")
|
||||
|
||||
remote_client = QdrantClient()
|
||||
prepare_collection(
|
||||
remote_client,
|
||||
COLLECTION_NAME,
|
||||
vectors_config={
|
||||
"all-minilm": models.VectorParams(size=384, distance=models.Distance.COSINE)
|
||||
},
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
remote_client.upsert(
|
||||
COLLECTION_NAME,
|
||||
points=[
|
||||
models.PointStruct(
|
||||
id=1,
|
||||
vector={
|
||||
"all-minilm": models.Document(
|
||||
text="bad text", model="sentence-transformers/all-MiniLM-L6-v2"
|
||||
),
|
||||
"bm25": models.Document(text="good text", model=MODEL_NAME),
|
||||
},
|
||||
)
|
||||
],
|
||||
)
|
||||
Reference in New Issue
Block a user