/
vector_search.py
54 lines (47 loc) · 1.37 KB
/
vector_search.py
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
from typing import Dict, Union, Callable
from deeplake.core.dataset import Dataset as DeepLakeDataset
from deeplake.core.vectorstore.vector_search import utils
from deeplake.core.vectorstore.vector_search import filter as filter_utils
from deeplake.core import vectorstore
def vector_search(
query,
query_emb,
exec_option,
dataset,
logger,
filter,
embedding_tensor,
distance_metric,
k,
return_tensors,
return_view,
token,
org_id,
return_tql,
) -> Union[Dict, DeepLakeDataset]:
runtime = utils.get_runtime_from_exec_option(exec_option)
if callable(filter):
raise ValueError(
f"UDF filter functions are not supported with the current `exec_option`={exec_option}. "
)
utils.check_indra_installation(exec_option)
view, tql_filter = filter_utils.attribute_based_filtering_tql(
view=dataset,
filter=filter,
logger=logger,
)
return vectorstore.indra_search_algorithm(
query_embedding=query_emb,
distance_metric=distance_metric.lower(),
deeplake_dataset=view,
k=k,
tql_string=query,
tql_filter=tql_filter,
embedding_tensor=embedding_tensor,
runtime=runtime,
return_tensors=return_tensors,
return_view=return_view,
token=token,
org_id=org_id,
return_tql=return_tql,
)