Skip to content
This repository was archived by the owner on Sep 9, 2026. It is now read-only.

Commit c7dfb90

Browse files
JohannesMessnerJohannesMessner
authored
feat: find function (#931)
* refactor: make abstract tensor parametrization modular Signed-off-by: Johannes Messner <[email protected]> * feat: add type for torch embedding Signed-off-by: Johannes Messner <[email protected]> * feat: embedding type for ndarray Signed-off-by: Johannes Messner <[email protected]> * fix: fix general embedding type Signed-off-by: Johannes Messner <[email protected]> * test: update tests Signed-off-by: Johannes Messner <[email protected]> * feat: find function Signed-off-by: Johannes Messner <[email protected]> * feat: batched query input Signed-off-by: Johannes Messner <[email protected]> * fix: fix metrics and add tests for them Signed-off-by: Johannes Messner <[email protected]> * test: add tests for find Signed-off-by: Johannes Messner <[email protected]> * test: add test for topk Signed-off-by: Johannes Messner <[email protected]> * feat: add batched find Signed-off-by: Johannes Messner <[email protected]> * test: add more tests Signed-off-by: Johannes Messner <[email protected]> * docs: improve docstrings Signed-off-by: Johannes Messner <[email protected]> * docs: improve docstrings Signed-off-by: Johannes Messner <[email protected]> * fix: mypy and some comments Signed-off-by: Johannes Messner <[email protected]> * refactor: address code review Signed-off-by: Johannes Messner <[email protected]> * fix: suppress mypy Signed-off-by: Johannes Messner <[email protected]> Signed-off-by: Johannes Messner <[email protected]>
1 parent dedd18e commit c7dfb90

14 files changed

Lines changed: 1132 additions & 0 deletions

File tree

‎docarray/typing/tensor/__init__.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from typing import Dict, Set, Type
2+
13
from docarray.typing.tensor.embedding import Embedding, NdArrayEmbedding
24
from docarray.typing.tensor.ndarray import NdArray
35
from docarray.typing.tensor.tensor import Tensor
@@ -7,8 +9,12 @@
79
'Tensor',
810
'Embedding',
911
'NdArrayEmbedding',
12+
'framework_types',
13+
'type_to_framework',
1014
]
1115

16+
framework_types: Dict[str, Set] = {'numpy': {NdArray, NdArrayEmbedding}, 'torch': set()}
17+
1218
try:
1319
import torch # noqa: F401
1420
except ImportError:
@@ -18,3 +24,10 @@
1824
from docarray.typing.tensor.torch_tensor import TorchTensor # noqa: F401
1925

2026
__all__.extend(['TorchEmbedding', 'TorchTensor'])
27+
framework_types['torch'] = {TorchTensor, TorchEmbedding}
28+
29+
type_to_framework: Dict[Type, str] = {
30+
type_: framework
31+
for framework, type_set in framework_types.items()
32+
for type_ in type_set
33+
}

‎docarray/utility/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from docarray.utility.find import find, find_batched
2+
3+
__all__ = ['find', 'find_batched']

‎docarray/utility/find.py‎

Lines changed: 321 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,321 @@
1+
import importlib
2+
from typing import Callable, List, NamedTuple, Optional, Type, Union
3+
4+
import torch # TODO(johannes) this breaks the optional import of torch
5+
6+
from docarray import Document, DocumentArray
7+
from docarray.typing import Tensor
8+
from docarray.typing.tensor import type_to_framework
9+
10+
# but will be fixed once we have a computational backend
11+
12+
13+
class FindResult(NamedTuple):
14+
documents: DocumentArray
15+
scores: Tensor
16+
17+
18+
def find(
19+
index: DocumentArray,
20+
query: Union[Tensor, Document],
21+
embedding_field: str = 'embedding',
22+
metric: str = 'cosine_sim',
23+
limit: int = 10,
24+
device: Optional[str] = None,
25+
descending: Optional[bool] = None,
26+
) -> FindResult:
27+
"""
28+
Find the closest Documents in the index to the query.
29+
Supports PyTorch and NumPy embeddings.
30+
31+
.. note::
32+
This utility function is likely to be removed once
33+
Document Stores are available.
34+
At that point, and in-memory Document Store will serve the same purpose
35+
by exposing a .find() method.
36+
37+
.. note::
38+
This is a simple implementation that assumes the same embedding field name for
39+
both query and index, does not support nested search, and does not support
40+
hybrid (multi-vector) search. These shortcoming will be addressed in future
41+
versions.
42+
43+
EXAMPLE USAGE
44+
45+
.. code-block:: python
46+
47+
from docarray import DocumentArray, Document
48+
from docarray.typing import TorchTensor
49+
from docarray.utility.find import find
50+
51+
52+
class MyDocument(Document):
53+
embedding: TorchTensor
54+
55+
56+
index = DocumentArray[MyDocument](
57+
[MyDocument(embedding=torch.rand(128)) for _ in range(100)]
58+
)
59+
60+
# use Document as query
61+
query = MyDocument(embedding=torch.rand(128))
62+
top_matches, scores = find(
63+
index=index,
64+
query=query,
65+
embedding_field='tensor',
66+
metric='cosine_sim',
67+
)
68+
69+
# use tensor as query
70+
query = torch.rand(128)
71+
top_matches, scores = find(
72+
index=index,
73+
query=query,
74+
embedding_field='tensor',
75+
metric='cosine_sim',
76+
)
77+
78+
:param index: the index of Documents to search in
79+
:param query: the query to search for
80+
:param embedding_field: the tensor-like field in the index to use
81+
for the similarity computation
82+
:param metric: the distance metric to use for the similarity computation.
83+
Can be one of the following strings:
84+
'cosine_sim' for cosine similarity, 'euclidean_dist' for euclidean distance,
85+
'sqeuclidean_dist' for squared euclidean distance
86+
:param limit: return the top `limit` results
87+
:param device: the computational device to use,
88+
can be either `cpu` or a `cuda` device.
89+
:param descending: sort the results in descending order.
90+
Per default, this is chosen based on the `metric` argument.
91+
:return: A named tuple of the form (DocumentArray, Tensor),
92+
where the first element contains the closes matches for the query,
93+
and the second element contains the corresponding scores.
94+
"""
95+
query = _extract_embedding_single(query, embedding_field)
96+
return find_batched(
97+
index=index,
98+
query=query,
99+
embedding_field=embedding_field,
100+
metric=metric,
101+
limit=limit,
102+
device=device,
103+
descending=descending,
104+
)[0]
105+
106+
107+
def find_batched(
108+
index: DocumentArray,
109+
query: Union[Tensor, DocumentArray],
110+
embedding_field: str = 'embedding',
111+
metric: str = 'cosine_sim',
112+
limit: int = 10,
113+
device: Optional[str] = None,
114+
descending: Optional[bool] = None,
115+
) -> List[FindResult]:
116+
"""
117+
Find the closest Documents in the index to the queries.
118+
Supports PyTorch and NumPy embeddings.
119+
120+
.. note::
121+
This utility function is likely to be removed once
122+
Document Stores are available.
123+
At that point, and in-memory Document Store will serve the same purpose
124+
by exposing a .find() method.
125+
126+
.. note::
127+
This is a simple implementation that assumes the same embedding field name for
128+
both query and index, does not support nested search, and does not support
129+
hybrid (multi-vector) search. These shortcoming will be addressed in future
130+
versions.
131+
132+
EXAMPLE USAGE
133+
134+
.. code-block:: python
135+
136+
from docarray import DocumentArray, Document
137+
from docarray.typing import TorchTensor
138+
from docarray.utility.find import find
139+
140+
141+
class MyDocument(Document):
142+
embedding: TorchTensor
143+
144+
145+
index = DocumentArray[MyDocument](
146+
[MyDocument(embedding=torch.rand(128)) for _ in range(100)]
147+
)
148+
149+
# use DocumentArray as query
150+
query = DocumentArray[MyDocument](
151+
[MyDocument(embedding=torch.rand(128)) for _ in range(3)]
152+
)
153+
results = find(
154+
index=index,
155+
query=query,
156+
embedding_field='tensor',
157+
metric='cosine_sim',
158+
)
159+
top_matches, scores = results[0]
160+
161+
# use tensor as query
162+
query = torch.rand(3, 128)
163+
results, scores = find(
164+
index=index,
165+
query=query,
166+
embedding_field='tensor',
167+
metric='cosine_sim',
168+
)
169+
top_matches, scores = results[0]
170+
171+
:param index: the index of Documents to search in
172+
:param query: the query to search for
173+
:param embedding_field: the tensor-like field in the index to use
174+
for the similarity computation
175+
:param metric: the distance metric to use for the similarity computation.
176+
Can be one of the following strings:
177+
'cosine_sim' for cosine similarity, 'euclidean_dist' for euclidean distance,
178+
'sqeuclidean_dist' for squared euclidean distance
179+
:param limit: return the top `limit` results
180+
:param device: the computational device to use,
181+
can be either `cpu` or a `cuda` device.
182+
:param descending: sort the results in descending order.
183+
Per default, this is chosen based on the `metric` argument.
184+
:return: a list of named tuples of the form (DocumentArray, Tensor),
185+
where the first element contains the closes matches for each query,
186+
and the second element contains the corresponding scores.
187+
"""
188+
if descending is None:
189+
descending = metric.endswith('_sim') # similarity metrics are descending
190+
191+
embedding_type = _da_attr_type(index, embedding_field)
192+
193+
# get framework-specific distance and top_k function
194+
metric_fn = _get_metric_fn(embedding_type, metric)
195+
top_k_fn = _get_topk_fn(embedding_type)
196+
197+
# extract embeddings from query and index
198+
index_embeddings = _extraxt_embeddings(index, embedding_field, embedding_type)
199+
query_embeddings = _extraxt_embeddings(query, embedding_field, embedding_type)
200+
201+
# compute distances and return top results
202+
dists = metric_fn(query_embeddings, index_embeddings, device=device)
203+
top_scores, top_indices = top_k_fn(
204+
dists, k=limit, device=device, descending=descending
205+
)
206+
207+
index_doc_type = index.document_type
208+
results = []
209+
for indices_per_query, scores_per_query in zip(top_indices, top_scores):
210+
docs_per_query = DocumentArray[index_doc_type]([]) # type: ignore
211+
for idx in indices_per_query: # workaround until #930 is fixed
212+
docs_per_query.append(index[idx])
213+
results.append(FindResult(scores=scores_per_query, documents=docs_per_query))
214+
return results
215+
216+
217+
def _extract_embedding_single(
218+
data: Union[DocumentArray, Document, Tensor],
219+
embedding_field: str,
220+
) -> Tensor:
221+
"""Extract the embeddings from a single query,
222+
and return it in a batched representation.
223+
224+
:param data: the data
225+
:param embedding_field: the embedding field
226+
:param embedding_type: type of the embedding: torch.Tensor, numpy.ndarray etc.
227+
:return: the embeddings
228+
"""
229+
if isinstance(data, Document):
230+
emb = getattr(data, embedding_field)
231+
else: # treat data as tensor
232+
emb = data
233+
if len(emb.shape) == 1:
234+
# TODO(johannes) solve this with computational backend,
235+
# this is ugly hack for now
236+
if isinstance(emb, torch.Tensor):
237+
emb = emb.unsqueeze(0)
238+
else:
239+
import numpy as np
240+
241+
if isinstance(emb, np.ndarray):
242+
emb = np.expand_dims(emb, axis=0)
243+
return emb
244+
245+
246+
def _extraxt_embeddings(
247+
data: Union[DocumentArray, Document, Tensor],
248+
embedding_field: str,
249+
embedding_type: Type,
250+
) -> Tensor:
251+
"""Extract the embeddings from the data.
252+
253+
:param data: the data
254+
:param embedding_field: the embedding field
255+
:param embedding_type: type of the embedding: torch.Tensor, numpy.ndarray etc.
256+
:return: the embeddings
257+
"""
258+
# TODO(johannes) put docarray stack in the computational backend
259+
if isinstance(data, DocumentArray):
260+
emb = getattr(data, embedding_field)
261+
if not data.is_stacked():
262+
emb = embedding_type.__docarray_stack__(emb)
263+
elif isinstance(data, Document):
264+
emb = getattr(data, embedding_field)
265+
else: # treat data as tensor
266+
emb = data
267+
268+
if len(emb.shape) == 1:
269+
# TODO(johannes) solve this with computational backend,
270+
# this is ugly hack for now
271+
if isinstance(emb, torch.Tensor):
272+
emb = emb.unsqueeze(0)
273+
else:
274+
import numpy as np
275+
276+
if isinstance(emb, np.ndarray):
277+
emb = np.expand_dims(emb, axis=0)
278+
return emb
279+
280+
281+
def _da_attr_type(da: DocumentArray, attr: str) -> Type:
282+
"""Get the type of the attribute according to the Document type
283+
(schema) of the DocumentArray.
284+
285+
:param da: the DocumentArray
286+
:param attr: the attribute name
287+
:return: the type of the attribute
288+
"""
289+
return da.document_type.__fields__[attr].type_
290+
291+
292+
def _get_topk_fn(embedding_type: Type) -> Callable:
293+
"""Dynamically import the distance function from the framework-specific module.
294+
This will go away once we have a computational backend.
295+
296+
:param embedding_type: the type of the embedding
297+
:param distance_name: the name of the distance function
298+
:return: the framework-specific distance function
299+
"""
300+
framework = type_to_framework[embedding_type]
301+
return getattr(
302+
importlib.import_module(f'docarray.utility.helper.{framework}'),
303+
'top_k',
304+
)
305+
306+
307+
def _get_metric_fn(embedding_type: Type, metric: Union[str, Callable]) -> Callable:
308+
"""Dynamically import the distance function from the framework-specific module.
309+
This will go away once we have a proper computational backend.
310+
311+
:param embedding_type: the type of the embedding
312+
:param metric: the name of the metric, or the metric itself
313+
:return: the framework-specific metric
314+
"""
315+
if callable(metric):
316+
return metric
317+
framework = type_to_framework[embedding_type]
318+
return getattr(
319+
importlib.import_module(f'docarray.utility.math.metrics.{framework}'),
320+
f'{metric}',
321+
)

‎docarray/utility/helper/__init__.py‎

Whitespace-only changes.

0 commit comments

Comments
 (0)