Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Local inference upload collection and upload records #862

Merged
merged 20 commits into from
Jan 29, 2025
Merged

Local inference upload collection and upload records #862

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
d0d2839
wip: add draft implementation of batch processing
joein Nov 22, 2024
4912805
fix: embed dict and list of docs, remove redundant code
joein Nov 23, 2024
e376f14
new: regen async, small refactor
joein Dec 4, 2024
be8f5d2
refactor: add docstrings, rename methods
joein Dec 5, 2024
b0dea66
Upload points local inference (#881)
joein Jan 15, 2025
4802a0f
refactor: refactor semi ordered map
joein Jan 20, 2025
718479e
tests: add test for local inference with batches with docs and vectors
joein Jan 21, 2025
dbec8fb
tests: check the order of dict processing
joein Jan 21, 2025
fed2f9c
new: distinguish models by options
joein Jan 22, 2025
7af5095
fix: fix typing
joein Jan 22, 2025
d808203
fix: fix types
joein Jan 22, 2025
caf1bb7
new: embed batches with different options
joein Jan 23, 2025
8090eb7
tests: add tests for batch with different options
joein Jan 23, 2025
1bf2d7e
fix: ignore ide incorrect type inspection
joein Jan 23, 2025
0dc8e23
tests: wait for points to be inserted
joein Jan 23, 2025
42a5172
fix: set threads to 1 in parallel inference
joein Jan 23, 2025
7976b88
new: adjust max internal batch size
joein Jan 23, 2025
01894a8
fix: fix type hints
joein Jan 23, 2025
48907e6
function to get embeddings size (#892)
generall Jan 29, 2025
0e3ae86
new: add local inference batch size (#894)
joein Jan 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
134 changes: 120 additions & 14 deletions qdrant_client/async_qdrant_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import warnings
from copy import deepcopy
from typing import Any, Awaitable, Callable, Iterable, Mapping, Optional, Sequence, Union
import numpy as np
from qdrant_client import grpc as grpc
from qdrant_client.async_client_base import AsyncQdrantBase
from qdrant_client.common.client_warnings import show_warning_once
Expand Down Expand Up @@ -95,6 +96,7 @@ def __init__(
Union[Callable[[], str], Callable[[], Awaitable[str]]]
] = None,
cloud_inference: bool = False,
local_inference_batch_size: Optional[int] = None,
check_compatibility: bool = True,
**kwargs: Any,
):
Expand Down Expand Up @@ -142,6 +144,7 @@ def __init__(
"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

async def close(self, grpc_grace: Optional[float] = None, **kwargs: Any) -> None:
"""Closes the connection to Qdrant
Expand Down Expand Up @@ -395,7 +398,11 @@ async def query_batch_points(
requests = self._resolve_query_batch_request(requests)
requires_inference = self._inference_inspector.inspect(requests)
if requires_inference and (not self.cloud_inference):
requests = [self._embed_models(request) for request in requests]
requests = list(
self._embed_models(
requests, is_query=True, batch_size=self.local_inference_batch_size
)
)
return await self._client.query_batch_points(
collection_name=collection_name,
requests=requests,
Expand Down Expand Up @@ -522,10 +529,35 @@ async def query_points(
query = self._resolve_query(query)
requires_inference = self._inference_inspector.inspect([query, prefetch])
if requires_inference and (not self.cloud_inference):
query = self._embed_models(query, is_query=True) if query is not None else None
prefetch = (
self._embed_models(prefetch, is_query=True) if prefetch is not None else None
query = (
next(
iter(
self._embed_models(
query, is_query=True, batch_size=self.local_inference_batch_size
)
)
)
if query is not None
else None
)
if isinstance(prefetch, list):
prefetch = list(
self._embed_models(
prefetch, is_query=True, batch_size=self.local_inference_batch_size
)
)
else:
prefetch = (
next(
iter(
self._embed_models(
prefetch, is_query=True, batch_size=self.local_inference_batch_size
)
)
)
if prefetch is not None
else None
)
return await self._client.query_points(
collection_name=collection_name,
query=query,
Expand Down Expand Up @@ -661,10 +693,31 @@ async def query_points_groups(
query = self._resolve_query(query)
requires_inference = self._inference_inspector.inspect([query, prefetch])
if requires_inference and (not self.cloud_inference):
query = self._embed_models(query, is_query=True) if query is not None else None
prefetch = (
self._embed_models(prefetch, is_query=True) if prefetch is not None else None
query = (
next(
iter(
self._embed_models(
query, is_query=True, batch_size=self.local_inference_batch_size
)
)
)
if query is not None
else None
)
if isinstance(prefetch, list):
prefetch = list(
self._embed_models(
prefetch, is_query=True, batch_size=self.local_inference_batch_size
)
)
elif prefetch is not None:
prefetch = next(
iter(
self._embed_models(
prefetch, is_query=True, batch_size=self.local_inference_batch_size
)
)
)
return await self._client.query_points_groups(
collection_name=collection_name,
query=query,
Expand Down Expand Up @@ -1506,10 +1559,20 @@ async def upsert(
)
requires_inference = self._inference_inspector.inspect(points)
if requires_inference and (not self.cloud_inference):
if isinstance(points, list):
points = [self._embed_models(point, is_query=False) for point in points]
if isinstance(points, types.Batch):
points = next(
iter(
self._embed_models(
points, is_query=False, batch_size=self.local_inference_batch_size
)
)
)
else:
points = self._embed_models(points, is_query=False)
points = list(
self._embed_models(
points, is_query=False, batch_size=self.local_inference_batch_size
)
)
return await self._client.upsert(
collection_name=collection_name,
points=points,
Expand Down Expand Up @@ -1560,7 +1623,11 @@ async def update_vectors(
assert len(kwargs) == 0, f"Unknown arguments: {list(kwargs.keys())}"
requires_inference = self._inference_inspector.inspect(points)
if requires_inference and (not self.cloud_inference):
points = [self._embed_models(point, is_query=False) for point in points]
points = list(
self._embed_models(
points, is_query=False, batch_size=self.local_inference_batch_size
)
)
return await self._client.update_vectors(
collection_name=collection_name,
points=points,
Expand Down Expand Up @@ -2000,9 +2067,11 @@ async def batch_update_points(
assert len(kwargs) == 0, f"Unknown arguments: {list(kwargs.keys())}"
requires_inference = self._inference_inspector.inspect(update_operations)
if requires_inference and (not self.cloud_inference):
update_operations = [
self._embed_models(op, is_query=False) for op in update_operations
]
update_operations = list(
self._embed_models(
update_operations, is_query=False, batch_size=self.local_inference_batch_size
)
)
return await self._client.batch_update_points(
collection_name=collection_name,
update_operations=update_operations,
Expand Down Expand Up @@ -2426,7 +2495,25 @@ def upload_points(
This parameter overwrites shard keys written in the records.

"""

def chain(*iterables: Iterable) -> Iterable:
for iterable in iterables:
yield from iterable

assert len(kwargs) == 0, f"Unknown arguments: {list(kwargs.keys())}"
if not self.cloud_inference:
iter_points = iter(points)
requires_inference = False
try:
point = next(iter_points)
requires_inference = self._inference_inspector.inspect(point)
points = chain(iter([point]), iter_points)
except (StopIteration, StopAsyncIteration):
points = []
if requires_inference:
points = self._embed_models_strict(
points, parallel=parallel, batch_size=self.local_inference_batch_size
)
return self._client.upload_points(
collection_name=collection_name,
points=points,
Expand Down Expand Up @@ -2478,7 +2565,26 @@ def upload_collection(
If multiple shard_keys are provided, the update will be written to each of them.
Only works for collections with `custom` sharding method.
"""

def chain(*iterables: Iterable) -> Iterable:
for iterable in iterables:
yield from iterable

assert len(kwargs) == 0, f"Unknown arguments: {list(kwargs.keys())}"
if not self.cloud_inference:
if not isinstance(vectors, dict) and (not isinstance(vectors, np.ndarray)):
requires_inference = False
try:
iter_vectors = iter(vectors)
vector = next(iter_vectors)
requires_inference = self._inference_inspector.inspect(vector)
vectors = chain(iter([vector]), iter_vectors)
except (StopIteration, StopAsyncIteration):
vectors = []
if requires_inference:
vectors = self._embed_models_strict(
vectors, parallel=parallel, batch_size=self.local_inference_batch_size
)
return self._client.upload_collection(
collection_name=collection_name,
vectors=vectors,
Expand Down
Loading