-
Notifications
You must be signed in to change notification settings - Fork 243
Expand file tree
/
Copy pathfind.py
More file actions
56 lines (51 loc) · 1.84 KB
/
Copy pathfind.py
File metadata and controls
56 lines (51 loc) · 1.84 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
from typing import TYPE_CHECKING, TypeVar, List, Union, Optional, Dict, Sequence
if TYPE_CHECKING:
import numpy as np
import tensorflow
import torch
# Define the expected input type that your ANN search supports
MilvusArrayType = TypeVar(
'MilvusArrayType',
np.ndarray,
tensorflow.Tensor,
torch.Tensor,
Sequence[float],
)
from docarray import Document, DocumentArray
class FindMixin:
def _find(
self,
query: 'MilvusArrayType',
limit: int = 10,
filter: Optional[Dict] = None,
param=None,
**kwargs
) -> List['DocumentArray']:
"""Returns `limit` approximate nearest neighbors given a batch of input queries.
If the query is a single query, should return a DocumentArray, otherwise a list of DocumentArrays containing
the closest Documents for each query.
"""
if param is None:
param = dict()
kwargs = self._update_kwargs_from_config('consistency_level', **kwargs)
with self.loaded_collection():
results = self._collection.search(
data=query,
anns_field='embedding',
limit=limit,
expr=filter,
param=param,
output_fields=['serialized'],
**kwargs,
)
return self._docs_from_search_response(results, distance=self._config.distance)
def _filter(self, filter, limit=10, **kwargs):
kwargs = self._update_kwargs_from_config('consistency_level', **kwargs)
with self.loaded_collection():
results = self._collection.query(
expr=filter,
limit=limit,
output_fields=['serialized'],
**kwargs,
)
return self._docs_from_query_response(results)[:limit]