from docarray import DocumentArray, Document
from docarray.utility import find, find_batched
from docarray.typing import TorchTensor

class MyDoc(Document):
    tensor: TorchTensor

da = DocumentArray[MyDoc](MyDoc(tensor=torch.rand(128)) for _ in range(10))

matches, scores = find(da, MyDoc(tensor=torch.rand(128)), embedding_field='tensor', metric='cosine_sim')

batched_query = DocumentArray[MyDoc](
    [MyDoc(tensor=torch.rand(128)) for _ in range(3)]
)
results = find_batched(da, batched_query, embedding_field='tensor', metric='cosine_sim')
assert len(results) == 3
for matches_i, socres_i in results:
    ...