From 1addc114cd3ae55f974e0a879b81588dec658d24 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 29 Nov 2022 11:13:57 +0100 Subject: [PATCH 01/17] refactor: make abstract tensor parametrization modular Signed-off-by: Johannes Messner --- docarray/typing/tensor/abstract_tensor.py | 46 +++++++++++++++++------ 1 file changed, 35 insertions(+), 11 deletions(-) diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py index 117e8fab4cb..208278a5264 100644 --- a/docarray/typing/tensor/abstract_tensor.py +++ b/docarray/typing/tensor/abstract_tensor.py @@ -20,8 +20,12 @@ class AbstractTensor(AbstractType, Generic[ShapeT], ABC): @abc.abstractmethod def __validate_shape__(cls, t: T, shape: Tuple[int]) -> T: """Every tensor has to implement this method in order to - enbale syntax of the form Tensor[shape]. - The intended behavoiour is as follows: + enable syntax of the form Tensor[shape]. + + It is called at "data setting time", + i.e. when a tensor is passed to a Document field of type Tensor[shape]. + + The intended behaviour is as follows: - If the shape of `t` is equal to `shape`, return `t`. - If the shape of `t` is not equal to `shape`, but can be reshaped to `shape`, return `t` reshaped to `shape`. @@ -34,6 +38,32 @@ def __validate_shape__(cls, t: T, shape: Tuple[int]) -> T: """ ... + @classmethod + def __validate_getitem__(cls, item: Any) -> Tuple[int]: + """This method validates the input to __class_getitem__. + + It is called at "class creation time", + i.e. when a class is created with syntax of the form Tensor[shape]. + + The default implementation tries to cast any `item` to a tuple of ints. + A subclass can override this method to implement custom validation logic. + + The output of this is eventually passed to + {ref}`AbstractTensor.__validate_shape__` as its `shape` argument. + + Raises `ValueError` if the input `item` does not pass validation. + + :param item: The item to validate, passed to __class_getitem__ (`Tensor[item]`). + :return: The validated item == the target shape of this tensor. + """ + if isinstance(item, int): + item = (item,) + try: + item = tuple(item) + except TypeError: + raise TypeError(f'{item} is not a valid tensor shape.') + return item + @classmethod def _create_parametrized_type(cls: Type[T], shape: Tuple[int]): shape_str = ', '.join([str(s) for s in shape]) @@ -58,12 +88,6 @@ def validate( return _ParametrizedTensor - def __class_getitem__(cls, item): - if isinstance(item, int): - item = (item,) - try: - item = tuple(item) - except TypeError: - raise TypeError(f'{item} is not a valid tensor shape.') - - return cls._create_parametrized_type(item) + def __class_getitem__(cls, item: Any): + target_shape = cls.__validate_getitem__(item) + return cls._create_parametrized_type(target_shape) From b0b8497baba1efc15ea7c82e0d0ef6282b3fa989 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 29 Nov 2022 11:24:19 +0100 Subject: [PATCH 02/17] feat: add type for torch embedding Signed-off-by: Johannes Messner --- docarray/typing/__init__.py | 3 ++- docarray/typing/tensor/__init__.py | 4 ++-- docarray/typing/tensor/torch_tensor.py | 12 ++++++++++++ tests/integrations/typing/test_torch_tensor.py | 10 ++++++++-- tests/units/typing/tensor/test_torch_tensor.py | 18 +++++++++++++++++- 5 files changed, 41 insertions(+), 6 deletions(-) diff --git a/docarray/typing/__init__.py b/docarray/typing/__init__.py index 581c7ff4bd0..f2f40172000 100644 --- a/docarray/typing/__init__.py +++ b/docarray/typing/__init__.py @@ -1,5 +1,5 @@ from docarray.typing.id import ID -from docarray.typing.tensor import NdArray, Tensor, TorchTensor +from docarray.typing.tensor import NdArray, Tensor, TorchEmbedding, TorchTensor from docarray.typing.tensor.embedding import Embedding from docarray.typing.url import AnyUrl, ImageUrl, TextUrl @@ -12,4 +12,5 @@ 'AnyUrl', 'ID', 'Tensor', + 'TorchEmbedding', ] diff --git a/docarray/typing/tensor/__init__.py b/docarray/typing/tensor/__init__.py index 5bb6553a9b8..d0de64a42bd 100644 --- a/docarray/typing/tensor/__init__.py +++ b/docarray/typing/tensor/__init__.py @@ -1,5 +1,5 @@ from docarray.typing.tensor.ndarray import NdArray from docarray.typing.tensor.tensor import Tensor -from docarray.typing.tensor.torch_tensor import TorchTensor +from docarray.typing.tensor.torch_tensor import TorchEmbedding, TorchTensor -__all__ = ['NdArray', 'TorchTensor', 'Tensor'] +__all__ = ['NdArray', 'TorchTensor', 'Tensor', 'TorchEmbedding'] diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py index b564908be75..52b0c84c6aa 100644 --- a/docarray/typing/tensor/torch_tensor.py +++ b/docarray/typing/tensor/torch_tensor.py @@ -212,3 +212,15 @@ def _flush_tensor_to_proto(pb_msg: 'NdArrayProto', value: 'TorchTensor'): pb_msg.dense.ClearField('shape') pb_msg.dense.shape.extend(list(value_np.shape)) pb_msg.dense.dtype = value_np.dtype.str + + +class TorchEmbedding(TorchTensor): + @classmethod + def __validate_getitem__(cls, item: Any) -> Tuple[int]: + shape = super().__validate_getitem__(item) + if len(shape) > 1: + raise ValueError( + '`TorchEmbedding` can only have a single dimension/axis.' + 'You can use `TorchTensor` as a general tensor instead.' + ) + return shape diff --git a/tests/integrations/typing/test_torch_tensor.py b/tests/integrations/typing/test_torch_tensor.py index 8c41ca62ad0..1e82b711b8d 100644 --- a/tests/integrations/typing/test_torch_tensor.py +++ b/tests/integrations/typing/test_torch_tensor.py @@ -1,15 +1,21 @@ import torch from docarray import Document -from docarray.typing import TorchTensor +from docarray.typing import TorchEmbedding, TorchTensor def test_set_torch_tensor(): class MyDocument(Document): tensor: TorchTensor + embedding: TorchEmbedding - d = MyDocument(tensor=torch.zeros((3, 224, 224))) + d = MyDocument(tensor=torch.zeros((3, 224, 224)), embedding=torch.zeros((128,))) assert isinstance(d.tensor, TorchTensor) assert isinstance(d.tensor, torch.Tensor) assert (d.tensor == torch.zeros((3, 224, 224))).all() + + assert isinstance(d.embedding, TorchTensor) + assert isinstance(d.embedding, TorchEmbedding) + assert isinstance(d.embedding, torch.Tensor) + assert (d.embedding == torch.zeros((128,))).all() diff --git a/tests/units/typing/tensor/test_torch_tensor.py b/tests/units/typing/tensor/test_torch_tensor.py index ac064d5489d..6e3b0e8f3de 100644 --- a/tests/units/typing/tensor/test_torch_tensor.py +++ b/tests/units/typing/tensor/test_torch_tensor.py @@ -3,7 +3,7 @@ from pydantic.tools import parse_obj_as, schema_json_of from docarray.document.io.json import orjson_dumps -from docarray.typing import TorchTensor +from docarray.typing import TorchEmbedding, TorchTensor def test_proto_tensor(): @@ -57,3 +57,19 @@ def test_parametrized(): # wrong and not reshapable shape with pytest.raises(ValueError): parse_obj_as(TorchTensor[3, 224, 224], torch.zeros(224, 224)) + + +def test_torch_embedding(): + # correct shape + tensor = parse_obj_as(TorchEmbedding[128], torch.zeros(128)) + assert isinstance(tensor, TorchEmbedding) + assert isinstance(tensor, torch.Tensor) + assert tensor.shape == (128,) + + # wrong shape at data setting time + with pytest.raises(ValueError): + parse_obj_as(TorchEmbedding[128], torch.zeros(256)) + + # illegal shape at class creation time + with pytest.raises(ValueError): + parse_obj_as(TorchEmbedding[128, 128], torch.zeros(128, 128)) From 68f3914da286a457bdbd164023f1cf3a07f4f02b Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 29 Nov 2022 11:58:29 +0100 Subject: [PATCH 03/17] feat: embedding type for ndarray Signed-off-by: Johannes Messner --- docarray/typing/tensor/__init__.py | 12 ++++++-- docarray/typing/tensor/embedding.py | 39 ++++++++++++++++-------- docarray/typing/tensor/torch_tensor.py | 12 -------- tests/units/typing/tensor/test_tensor.py | 18 +++++++++++ 4 files changed, 55 insertions(+), 26 deletions(-) diff --git a/docarray/typing/tensor/__init__.py b/docarray/typing/tensor/__init__.py index d0de64a42bd..b0d51577632 100644 --- a/docarray/typing/tensor/__init__.py +++ b/docarray/typing/tensor/__init__.py @@ -1,5 +1,13 @@ +from docarray.typing.tensor.embedding import Embedding, NdArrayEmbedding, TorchEmbedding from docarray.typing.tensor.ndarray import NdArray from docarray.typing.tensor.tensor import Tensor -from docarray.typing.tensor.torch_tensor import TorchEmbedding, TorchTensor +from docarray.typing.tensor.torch_tensor import TorchTensor -__all__ = ['NdArray', 'TorchTensor', 'Tensor', 'TorchEmbedding'] +__all__ = [ + 'NdArray', + 'TorchTensor', + 'Tensor', + 'Embedding', + 'TorchEmbedding', + 'NdArrayEmbedding', +] diff --git a/docarray/typing/tensor/embedding.py b/docarray/typing/tensor/embedding.py index a77db717098..35924369c13 100644 --- a/docarray/typing/tensor/embedding.py +++ b/docarray/typing/tensor/embedding.py @@ -1,18 +1,33 @@ -from typing import TypeVar +from abc import ABC +from typing import Any, Tuple, TypeVar, Union -from docarray.proto import NodeProto -from docarray.typing.tensor import NdArray +from docarray.typing.tensor.abstract_tensor import AbstractTensor +from docarray.typing.tensor.ndarray import NdArray +from docarray.typing.tensor.torch_tensor import TorchTensor T = TypeVar('T', bound='Embedding') -class Embedding(NdArray): - def _to_node_protobuf(self: T, field: str = 'tensor') -> NodeProto: - """Convert Document into a NodeProto protobuf message. This function should - be called when the Document is nested into another Document that need to be - converted into a protobuf - :param field: field in which to store the content in the node proto - :return: the nested item protobuf message - """ +class EmbeddingMixin(AbstractTensor, ABC): + alternative_type = None - return super()._to_node_protobuf(field='embedding') + @classmethod + def __validate_getitem__(cls, item: Any) -> Tuple[int]: + shape = super().__validate_getitem__(item) + if len(shape) > 1: + error_msg = f'`{cls}` can only have a single dimension/axis.' + if cls.alternative_type: + error_msg += f' Consider using {cls.alternative_type} instead.' + raise ValueError(error_msg) + return shape + + +class NdArrayEmbedding(NdArray, EmbeddingMixin): + alternative_type = NdArray + + +class TorchEmbedding(TorchTensor, EmbeddingMixin): + alternative_type = TorchTensor + + +Embedding = Union[TorchEmbedding, NdArrayEmbedding] diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py index 52b0c84c6aa..b564908be75 100644 --- a/docarray/typing/tensor/torch_tensor.py +++ b/docarray/typing/tensor/torch_tensor.py @@ -212,15 +212,3 @@ def _flush_tensor_to_proto(pb_msg: 'NdArrayProto', value: 'TorchTensor'): pb_msg.dense.ClearField('shape') pb_msg.dense.shape.extend(list(value_np.shape)) pb_msg.dense.dtype = value_np.dtype.str - - -class TorchEmbedding(TorchTensor): - @classmethod - def __validate_getitem__(cls, item: Any) -> Tuple[int]: - shape = super().__validate_getitem__(item) - if len(shape) > 1: - raise ValueError( - '`TorchEmbedding` can only have a single dimension/axis.' - 'You can use `TorchTensor` as a general tensor instead.' - ) - return shape diff --git a/tests/units/typing/tensor/test_tensor.py b/tests/units/typing/tensor/test_tensor.py index d0ef58bc581..dd8356ed84d 100644 --- a/tests/units/typing/tensor/test_tensor.py +++ b/tests/units/typing/tensor/test_tensor.py @@ -5,6 +5,7 @@ from docarray.document.io.json import orjson_dumps from docarray.typing import NdArray +from docarray.typing.tensor import NdArrayEmbedding def test_proto_tensor(): @@ -72,3 +73,20 @@ def test_parametrized(): # wrong and not reshapable shape with pytest.raises(ValueError): parse_obj_as(NdArray[3, 224, 224], np.zeros((224, 224))) + + +def test_np_embedding(): + # correct shape + tensor = parse_obj_as(NdArrayEmbedding[128], np.zeros((128,))) + assert isinstance(tensor, NdArrayEmbedding) + assert isinstance(tensor, NdArray) + assert isinstance(tensor, np.ndarray) + assert tensor.shape == (128,) + + # wrong shape at data setting time + with pytest.raises(ValueError): + parse_obj_as(NdArrayEmbedding[128], np.zeros((256,))) + + # illegal shape at class creation time + with pytest.raises(ValueError): + parse_obj_as(NdArrayEmbedding[128, 128], np.zeros((128, 128))) From 953b6908a90970462efd57192c358d833780264b Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 29 Nov 2022 12:03:19 +0100 Subject: [PATCH 04/17] fix: fix general embedding type Signed-off-by: Johannes Messner --- docarray/typing/tensor/embedding.py | 2 +- tests/integrations/document/test_proto.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/docarray/typing/tensor/embedding.py b/docarray/typing/tensor/embedding.py index 35924369c13..acf5b3ece16 100644 --- a/docarray/typing/tensor/embedding.py +++ b/docarray/typing/tensor/embedding.py @@ -30,4 +30,4 @@ class TorchEmbedding(TorchTensor, EmbeddingMixin): alternative_type = TorchTensor -Embedding = Union[TorchEmbedding, NdArrayEmbedding] +Embedding = Union[NdArrayEmbedding, TorchEmbedding] diff --git a/tests/integrations/document/test_proto.py b/tests/integrations/document/test_proto.py index 98f05f09e15..b790b3cfa3e 100644 --- a/tests/integrations/document/test_proto.py +++ b/tests/integrations/document/test_proto.py @@ -9,8 +9,10 @@ NdArray, Tensor, TextUrl, + TorchEmbedding, TorchTensor, ) +from docarray.typing.tensor import NdArrayEmbedding def test_multi_modal_doc_proto(): @@ -41,6 +43,8 @@ class MyDoc(Document): generic_nd_array: Tensor generic_torch_tensor: Tensor embedding: Embedding + torch_embedding: TorchEmbedding[128] + np_embedding: NdArrayEmbedding[128] doc = MyDoc( img_url='test.png', @@ -53,6 +57,8 @@ class MyDoc(Document): generic_nd_array=np.zeros((3, 224, 224)), generic_torch_tensor=torch.zeros((3, 224, 224)), embedding=np.zeros((3, 224, 224)), + torch_embedding=torch.zeros((128,)), + np_embedding=np.zeros((128,)), ) doc = MyDoc.from_protobuf(doc.to_protobuf()) @@ -78,4 +84,10 @@ class MyDoc(Document): assert (doc.generic_torch_tensor == torch.zeros((3, 224, 224))).all() assert isinstance(doc.generic_torch_tensor, torch.Tensor) + assert (doc.torch_embedding == torch.zeros((128,))).all() + assert isinstance(doc.torch_embedding, torch.Tensor) + + assert (doc.np_embedding == np.zeros((128,))).all() + assert isinstance(doc.np_embedding, np.ndarray) + assert (doc.embedding == np.zeros((3, 224, 224))).all() From 766f6b9ba1bdf4bf05c90c0105dd025236fc23b5 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Wed, 30 Nov 2022 15:17:46 +0100 Subject: [PATCH 05/17] test: update tests Signed-off-by: Johannes Messner --- tests/integrations/typing/test_embedding.py | 1 - tests/integrations/typing/test_typing_proto.py | 6 +++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/integrations/typing/test_embedding.py b/tests/integrations/typing/test_embedding.py index dbb5de3cca4..461f8be3714 100644 --- a/tests/integrations/typing/test_embedding.py +++ b/tests/integrations/typing/test_embedding.py @@ -10,6 +10,5 @@ class MyDocument(Document): d = MyDocument(embedding=np.zeros((3, 224, 224))) - assert isinstance(d.embedding, Embedding) assert isinstance(d.embedding, np.ndarray) assert (d.embedding == np.zeros((3, 224, 224))).all() diff --git a/tests/integrations/typing/test_typing_proto.py b/tests/integrations/typing/test_typing_proto.py index e1c839b5d48..a2b67e6a175 100644 --- a/tests/integrations/typing/test_typing_proto.py +++ b/tests/integrations/typing/test_typing_proto.py @@ -27,4 +27,8 @@ class Mymmdoc(Document): new_doc = AnyDocument.from_protobuf(doc.to_protobuf()) for field, value in new_doc: - assert isinstance(value, doc._get_nested_document_class(field)) + if field == 'embedding': + # embedding is a Union type, not supported by isinstance + assert isinstance(value, np.ndarray) or isinstance(value, torch.Tensor) + else: + assert isinstance(value, doc._get_nested_document_class(field)) From cd7087b1aa1996616266a1328ed865eb19a5fe14 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Fri, 9 Dec 2022 14:28:13 +0100 Subject: [PATCH 06/17] feat: find function Signed-off-by: Johannes Messner --- docarray/typing/tensor/__init__.py | 11 +++ docarray/utility/__init__.py | 3 + docarray/utility/distance.py | 1 + docarray/utility/find.py | 106 ++++++++++++++++++++ docarray/utility/helper/__init__.py | 0 docarray/utility/helper/numpy.py | 45 +++++++++ docarray/utility/helper/torch.py | 29 ++++++ docarray/utility/math/__init__.py | 0 docarray/utility/math/distance/__init__.py | 0 docarray/utility/math/distance/numpy.py | 81 ++++++++++++++++ docarray/utility/math/distance/torch.py | 107 +++++++++++++++++++++ 11 files changed, 383 insertions(+) create mode 100644 docarray/utility/__init__.py create mode 100644 docarray/utility/distance.py create mode 100644 docarray/utility/find.py create mode 100644 docarray/utility/helper/__init__.py create mode 100644 docarray/utility/helper/numpy.py create mode 100644 docarray/utility/helper/torch.py create mode 100644 docarray/utility/math/__init__.py create mode 100644 docarray/utility/math/distance/__init__.py create mode 100644 docarray/utility/math/distance/numpy.py create mode 100644 docarray/utility/math/distance/torch.py diff --git a/docarray/typing/tensor/__init__.py b/docarray/typing/tensor/__init__.py index 58568a08c36..abb578ff401 100644 --- a/docarray/typing/tensor/__init__.py +++ b/docarray/typing/tensor/__init__.py @@ -7,8 +7,12 @@ 'Tensor', 'Embedding', 'NdArrayEmbedding', + 'framework_types', + 'type_to_framework', ] +framework_types = {'numpy': {NdArray, NdArrayEmbedding}, 'torch': set()} + try: import torch # noqa: F401 except ImportError: @@ -18,3 +22,10 @@ from docarray.typing.tensor.torch_tensor import TorchTensor # noqa: F401 __all__.extend(['TorchEmbedding', 'TorchTensor']) + framework_types['torch'] = {TorchTensor, TorchEmbedding} + +type_to_framework = { + type_: framework + for framework, type_set in framework_types.items() + for type_ in type_set +} diff --git a/docarray/utility/__init__.py b/docarray/utility/__init__.py new file mode 100644 index 00000000000..69148798f93 --- /dev/null +++ b/docarray/utility/__init__.py @@ -0,0 +1,3 @@ +from docarray.utility.find import find + +__all__ = ['find'] diff --git a/docarray/utility/distance.py b/docarray/utility/distance.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/docarray/utility/distance.py @@ -0,0 +1 @@ + diff --git a/docarray/utility/find.py b/docarray/utility/find.py new file mode 100644 index 00000000000..3bc2527ade8 --- /dev/null +++ b/docarray/utility/find.py @@ -0,0 +1,106 @@ +import importlib +from typing import Callable, NamedTuple, Optional, Type, Union + +from docarray import Document, DocumentArray +from docarray.typing import Tensor +from docarray.typing.tensor import type_to_framework + + +class FindResult(NamedTuple): + documents: DocumentArray + scores: Tensor + + +def find( + index: DocumentArray, + query: Union[Tensor, Document, DocumentArray], + embedding_field: Optional[str] = 'embedding', + metric: Union[str, Callable[['Tensor', 'Tensor'], 'Tensor']] = 'cosine', + limit: int = 10, + device: Optional[str] = None, +) -> FindResult: + """ + Find the closest Documents in the index to the query. + + .. note:: + This utility function is likely to be removed once + Document Stores are available. + At that point, and in-memory Document Store will serve the same purpose + by exposing a .find() method. + + .. note:: + This is a simple implementation that assumes the same embedding field name for + both query and index, does not support nested search, and does not support + hybrid (multi-vector) search. These shortcoming will be addressed in future + versions. + + :param index: the index of Documents to search in + :param query: the query to search for + :param embedding_field: the tensor-like field in the index to use + for the similarity computation + :param metric: the distance metric to use for the similarity computation. + Can be a string specifying a predefined metric + (TODO(johannes) specify) or a callable that takes two tensors and returns + a tensor of distances. TODO(johannes) implement passing a callable + :param limit: return the top `limit` results + :param device: the computational device to use, + can be either `cpu` or a `cuda` device. + :return: the closest Documents in the index to the query + """ + embedding_type = _da_attr_type(index, embedding_field) + + # get framework-specific distance and top_k function + distance_fn = _get_distance_fn(embedding_type, metric) + top_k_fn = _get_topk_fn(embedding_type) + + index_embeddings = getattr(index, embedding_field) + if not index.is_stacked(): + index_embeddings = embedding_type.__docarray_stack__(index_embeddings) + + dists = distance_fn( + index_embeddings, getattr(query, embedding_field), device=device + ) + top_scores, top_indices = top_k_fn(dists, k=limit, device=device) + results_docs = DocumentArray( + index[i] for i in top_indices + ) # workaround until #930 is fixed + return FindResult(documents=results_docs, scores=top_scores) + + +def _da_attr_type(da: DocumentArray, attr: str) -> Type: + """Get the type of the attribute according to the Document type + (schema) of the DocumentArray. + + :param da: the DocumentArray + :param attr: the attribute name + :return: the type of the attribute + """ + return da.document_type.__fields__[attr].type_ + + +def _get_topk_fn(embedding_type: Type) -> Callable: + """Dynamically import the distance function from the framework-specific module. + + :param embedding_type: the type of the embedding + :param distance_name: the name of the distance function + :return: the framework-specific distance function + """ + framework = type_to_framework[embedding_type] + return getattr( + importlib.import_module(f'docarray.utility.helper.{framework}'), + 'top_k', + ) + + +def _get_distance_fn(embedding_type: Type, distance_name: str) -> Callable: + """Dynamically import the distance function from the framework-specific module. + + :param embedding_type: the type of the embedding + :param distance_name: the name of the distance function + :return: the framework-specific distance function + """ + framework = type_to_framework[embedding_type] + return getattr( + importlib.import_module(f'docarray.utility.math.distance.{framework}'), + f'{distance_name}', + ) diff --git a/docarray/utility/helper/__init__.py b/docarray/utility/helper/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/docarray/utility/helper/numpy.py b/docarray/utility/helper/numpy.py new file mode 100644 index 00000000000..d4d41f72a03 --- /dev/null +++ b/docarray/utility/helper/numpy.py @@ -0,0 +1,45 @@ +import warnings +from typing import Optional, Tuple + +import numpy as np + + +def top_k( + values: 'np.ndarray', + k: int, + descending: bool = False, + device: Optional[str] = None, +) -> Tuple['np.ndarray', 'np.ndarray']: + """Retrieves the top k smallest values in `values`, + and returns them alongside their indices in the input `values`. + Can also be used to retrieve the top k largest values, + by setting the `descending` flag. + + :param values: Torch tensor of values to rank. + Should be a vector (single axis), or squeezable. + TODO(johannes) is squeezable a term understood in np land? + :param k: number of values to retrieve + :param descending: retrieve largest values instead of smallest values + :param device: Not supported for this backend + :return: Tuple containing the retrieved values, and their indices + """ + if device is not None: + warnings.warn('`device` is not supported for numpy operations') + + if descending: + values = -values + + if k >= values.shape[1]: + idx = values.argsort(axis=1)[:, :k] + values = np.take_along_axis(values, idx, axis=1) + else: + idx_ps = values.argpartition(kth=k, axis=1)[:, :k] + values = np.take_along_axis(values, idx_ps, axis=1) + idx_fs = values.argsort(axis=1) + idx = np.take_along_axis(idx_ps, idx_fs, axis=1) + values = np.take_along_axis(values, idx_fs, axis=1) + + if descending: + values = -values + + return values, idx diff --git a/docarray/utility/helper/torch.py b/docarray/utility/helper/torch.py new file mode 100644 index 00000000000..c1f058a804c --- /dev/null +++ b/docarray/utility/helper/torch.py @@ -0,0 +1,29 @@ +from typing import Optional, Tuple + +import torch + + +def top_k( + values: 'torch.tensor', + k: int, + descending: bool = False, + device: Optional[str] = None, +) -> Tuple['torch.tensor', 'torch.tensor']: + """Retrieves the top k smallest values in `values`, + and returns them alongside their indices in the input `values`. + Can also be used to retrieve the top k largest values, + by setting the `descending` flag. + + :param values: Torch tensor of values to rank. + Should be a vector (single axis), or squeezable. + :param k: number of values to retrieve + :param descending: retrieve largest values instead of smallest values + :param device: the computational device to use, + can be either `cpu` or a `cuda` device. + :return: Tuple containing the retrieved values, and their indices + """ + if device is not None: + values = values.to(device) + values = values.squeeze() + k = max(k, len(values)) + return torch.topk(input=values, k=k, largest=descending, sorted=True) diff --git a/docarray/utility/math/__init__.py b/docarray/utility/math/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/docarray/utility/math/distance/__init__.py b/docarray/utility/math/distance/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/docarray/utility/math/distance/numpy.py b/docarray/utility/math/distance/numpy.py new file mode 100644 index 00000000000..eea654925be --- /dev/null +++ b/docarray/utility/math/distance/numpy.py @@ -0,0 +1,81 @@ +import warnings +from typing import Optional + +import numpy as np + + +def cosine( + x_mat: 'np.ndarray', + y_mat: 'np.ndarray', + eps: float = 1e-7, + device: Optional[str] = None, +) -> 'np.ndarray': + """Pairwise cosine distances between all vectors in x_mat and y_mat. + + :param x_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is the + number of vectors and n_dim is the number of dimensions of each example. + :param y_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is the + number of vectors and n_dim is the number of dimensions of each example. + :param eps: a small jitter to avoid divde by zero + :param device: Not supported for this backend + :return: np.ndarray of shape (n_vectors, n_vectors) containing all pairwise + cosine distances. + The index [i_x, i_y] contains the cosine distance between + x_mat[i_x] and y_mat[i_y]. + """ + if device is not None: + warnings.warn('`device` is not supported for numpy operations') + return 1 - np.clip( + (np.dot(x_mat, y_mat.T) + eps) + / ( + np.outer(np.linalg.norm(x_mat, axis=1), np.linalg.norm(y_mat, axis=1)) + eps + ), + -1, + 1, + ) + + +def sqeuclidean( + x_mat: 'np.ndarray', y_mat: 'np.ndarray', device: Optional[str] = None +) -> 'np.ndarray': + """Pairwise Squared Euclidian distances between all vectors in x_mat and y_mat. + + :param x_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is the + number of vectors and n_dim is the number of dimensions of each example. + :param y_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is the + number of vectors and n_dim is the number of dimensions of each example. + :param eps: a small jitter to avoid divde by zero + :param device: Not supported for this backend + :return: np.ndarray of shape (n_vectors, n_vectors) containing all pairwise + Squared Euclidian distances. + The index [i_x, i_y] contains the cosine Squared Euclidian between + x_mat[i_x] and y_mat[i_y]. + """ + if device is not None: + warnings.warn('`device` is not supported for numpy operations') + return ( + np.sum(y_mat**2, axis=1) + + np.sum(x_mat**2, axis=1)[:, np.newaxis] + - 2 * np.dot(x_mat, y_mat.T) + ) + + +def euclidean( + x_mat: 'np.ndarray', y_mat: 'np.ndarray', device: Optional[str] = None +) -> 'np.ndarray': + """Pairwise Euclidian distances between all vectors in x_mat and y_mat. + + :param x_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is the + number of vectors and n_dim is the number of dimensions of each example. + :param y_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is the + number of vectors and n_dim is the number of dimensions of each example. + :param eps: a small jitter to avoid divde by zero + :param device: Not supported for this backend + :return: np.ndarray of shape (n_vectors, n_vectors) containing all pairwise + euclidian distances. + The index [i_x, i_y] contains the euclidian distance between + x_mat[i_x] and y_mat[i_y]. + """ + if device is not None: + warnings.warn('`device` is not supported for numpy operations') + return np.sqrt(sqeuclidean(x_mat, y_mat)) diff --git a/docarray/utility/math/distance/torch.py b/docarray/utility/math/distance/torch.py new file mode 100644 index 00000000000..4eba79e5441 --- /dev/null +++ b/docarray/utility/math/distance/torch.py @@ -0,0 +1,107 @@ +from typing import TYPE_CHECKING, List, Optional + +import torch + +if TYPE_CHECKING: # pragma: no cover + import numpy + from torch import tensor + + +def _unsqueeze_if_single_axis(*matrices: 'tensor') -> List['tensor']: + """Unsqueezes tensors that only have one axis, at dim 0. + This ensures that all outputs can be treated as matrices, not vectors. + + :param matrices: Matrices to be unsqueezed + :return: List of the input matrices, + where single axis matrices are unsqueezed at dim 0. + """ + unsqueezed = [] + for m in matrices: + if len(m.shape) == 1: + unsqueezed.append(m.unsqueeze(0)) + else: + unsqueezed.append(m) + return unsqueezed + + +def cosine( + x_mat: 'tensor', y_mat: 'tensor', eps: float = 1e-7, device: Optional[str] = None +) -> 'numpy.ndarray': + """Pairwise cosine distances between all vectors in x_mat and y_mat. + + :param x_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the + number of vectors and n_dim is the number of dimensions of each example. + :param y_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the + number of vectors and n_dim is the number of dimensions of each example. + :param eps: a small jitter to avoid divde by zero + :param device: the device to use for pytorch computations. + Either 'cpu' or a 'cuda' device. + If not provided, the devices of x_mat and y_mat are used. + :return: np.ndarray of shape (n_vectors, n_vectors) containing all pairwise + cosine distances. + The index [i_x, i_y] contains the cosine distance between + x_mat[i_x] and y_mat[i_y]. + """ + if device is not None: + x_mat = x_mat.to(device) + y_mat = y_mat.to(device) + + x_mat, y_mat = _unsqueeze_if_single_axis(x_mat, y_mat) + + a_n, b_n = x_mat.norm(dim=1)[:, None], y_mat.norm(dim=1)[:, None] + a_norm = x_mat / torch.clamp(a_n, min=eps) + b_norm = y_mat / torch.clamp(b_n, min=eps) + return 1 - torch.mm(a_norm, b_norm.transpose(0, 1)) + + +def euclidean( + x_mat: 'tensor', y_mat: 'tensor', device: Optional[str] = None +) -> 'tensor': + """Pairwise Euclidian distances between all vectors in x_mat and y_mat. + + :param x_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the + number of vectors and n_dim is the number of dimensions of each example. + :param y_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the + number of vectors and n_dim is the number of dimensions of each example. + :param device: the device to use for pytorch computations. + Either 'cpu' or a 'cuda' device. + If not provided, the devices of x_mat and y_mat are used. + :return: np.ndarray of shape (n_vectors, n_vectors) containing all pairwise + euclidian distances. + The index [i_x, i_y] contains the euclidian distance between + x_mat[i_x] and y_mat[i_y]. + """ + if device is not None: + x_mat = x_mat.to(device) + y_mat = y_mat.to(device) + + x_mat, y_mat = _unsqueeze_if_single_axis(x_mat, y_mat) + + return torch.cdist(x_mat, y_mat) + + +def sqeuclidean( + x_mat: 'tensor', y_mat: 'tensor', device: Optional[str] = None +) -> 'tensor': + """Pairwise Squared Euclidian distances between all vectors in x_mat and y_mat. + + :param x_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the + number of vectors and n_dim is the number of dimensions of each example. + :param y_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the + number of vectors and n_dim is the number of dimensions of each example. + :param eps: a small jitter to avoid divde by zero + :param device: the device to use for pytorch computations. + Either 'cpu' or a 'cuda' device. + If not provided, the devices of x_mat and y_mat are used. + :return: np.ndarray of shape (n_vectors, n_vectors) containing all pairwise + Squared Euclidian distances. + The index [i_x, i_y] contains the cosine Squared Euclidian between + x_mat[i_x] and y_mat[i_y]. + """ + if device is not None: + x_mat = x_mat.to(device) + y_mat = y_mat.to(device) + + x_mat, y_mat = _unsqueeze_if_single_axis(x_mat, y_mat) + + return torch.cdist(x_mat, y_mat) ** 2 From 9e950e1252c4da4f5edd9c4e9cbfb42b912b6715 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 12 Dec 2022 17:58:19 +0100 Subject: [PATCH 07/17] feat: batched query input Signed-off-by: Johannes Messner --- docarray/utility/find.py | 88 +++++++++++++++++++++---- docarray/utility/helper/torch.py | 7 +- docarray/utility/math/distance/numpy.py | 28 +++++++- 3 files changed, 106 insertions(+), 17 deletions(-) diff --git a/docarray/utility/find.py b/docarray/utility/find.py index 3bc2527ade8..ac56ecc16da 100644 --- a/docarray/utility/find.py +++ b/docarray/utility/find.py @@ -1,5 +1,7 @@ import importlib -from typing import Callable, NamedTuple, Optional, Type, Union +from typing import Callable, List, NamedTuple, Optional, Type, Union + +import torch from docarray import Document, DocumentArray from docarray.typing import Tensor @@ -7,7 +9,8 @@ class FindResult(NamedTuple): - documents: DocumentArray + # for single query | for multiple queries + documents: Union[DocumentArray, List[DocumentArray]] scores: Tensor @@ -53,18 +56,79 @@ def find( distance_fn = _get_distance_fn(embedding_type, metric) top_k_fn = _get_topk_fn(embedding_type) - index_embeddings = getattr(index, embedding_field) - if not index.is_stacked(): - index_embeddings = embedding_type.__docarray_stack__(index_embeddings) + # extract embeddings from query and index + index_embeddings = _extraxt_embeddings(index, embedding_field, embedding_type) + query_embeddings = _extraxt_embeddings(query, embedding_field, embedding_type) - dists = distance_fn( - index_embeddings, getattr(query, embedding_field), device=device - ) + # compute distances and return top results + dists = distance_fn(query_embeddings, index_embeddings, device=device) top_scores, top_indices = top_k_fn(dists, k=limit, device=device) - results_docs = DocumentArray( - index[i] for i in top_indices - ) # workaround until #930 is fixed - return FindResult(documents=results_docs, scores=top_scores) + + result_docs = [] + to_da = True + for top_idx in top_indices: # workaround until #930 is fixed + if len(top_idx) == 0: + result_docs.append(index[top_idx]) + else: + inner_result_docs = [] + to_da = False + for inner_top_idx in top_idx: + inner_result_docs.append(index[inner_top_idx]) + if to_da: + result_docs = DocumentArray(result_docs) + + return FindResult(documents=result_docs, scores=top_scores) + + +def _extraxt_embeddings( + data: Union[DocumentArray, Document, Tensor], + embedding_field: str, + embedding_type: Type, +) -> Tensor: + """Extract the embeddings from the data. + + :param data: the data + :param embedding_field: the embedding field + :param embedding_type: type of the embedding: torch.Tensor, numpy.ndarray etc. + :return: the embeddings + """ + # TODO(johannes) put docarray stack in the computational backend + if isinstance(data, DocumentArray): + emb = getattr(data, embedding_field) + if not data.is_stacked(): + emb = embedding_type.__docarray_stack__(emb) + elif isinstance(data, Document): + emb = getattr(data, embedding_field) + else: # treat data as tensor + emb = data + + if len(emb.shape) == 1: + # TODO(johannes) solve this with computational backend, + # this is ugly hack for now + if isinstance(emb, torch.Tensor): + emb = emb.unsqueeze(0) + else: + import numpy as np + + if isinstance(emb, np.ndarray): + emb = np.expand_dims(emb, axis=0) + return emb + + +def _to_documentarray_query( + query: Union[Tensor, Document, DocumentArray], embedding_field: str +) -> DocumentArray: + """Convert the query to a DocumentArray. + + :param query: the query + :return: the query as DocumentArray + """ + if isinstance(query, DocumentArray): + return query + elif isinstance(query, Document): + return DocumentArray([query]) + else: + return DocumentArray([Document(**{embedding_field: query})]) def _da_attr_type(da: DocumentArray, attr: str) -> Type: diff --git a/docarray/utility/helper/torch.py b/docarray/utility/helper/torch.py index c1f058a804c..f23a19c58a3 100644 --- a/docarray/utility/helper/torch.py +++ b/docarray/utility/helper/torch.py @@ -15,7 +15,7 @@ def top_k( by setting the `descending` flag. :param values: Torch tensor of values to rank. - Should be a vector (single axis), or squeezable. + Has to be of shape (n_queries, n_values_per_query) :param k: number of values to retrieve :param descending: retrieve largest values instead of smallest values :param device: the computational device to use, @@ -24,6 +24,5 @@ def top_k( """ if device is not None: values = values.to(device) - values = values.squeeze() - k = max(k, len(values)) - return torch.topk(input=values, k=k, largest=descending, sorted=True) + k = max(k, len(values[-1])) + return torch.topk(input=values, k=k, largest=descending, sorted=True, dim=-1) diff --git a/docarray/utility/math/distance/numpy.py b/docarray/utility/math/distance/numpy.py index eea654925be..5aae8a03aaa 100644 --- a/docarray/utility/math/distance/numpy.py +++ b/docarray/utility/math/distance/numpy.py @@ -1,9 +1,26 @@ import warnings -from typing import Optional +from typing import List, Optional import numpy as np +def _expand_if_single_axis(*matrices: 'np.array') -> List['np.array']: + """Expands arrays that only have one axis, at dim 0. + This ensures that all outputs can be treated as matrices, not vectors. + + :param matrices: Matrices to be expanded + :return: List of the input matrices, + where single axis matrices are expanded at dim 0. + """ + expanded = [] + for m in matrices: + if len(m.shape) == 1: + expanded.append(np.expand_dims(m, axis=0)) + else: + expanded.append(m) + return expanded + + def cosine( x_mat: 'np.ndarray', y_mat: 'np.ndarray', @@ -25,6 +42,9 @@ def cosine( """ if device is not None: warnings.warn('`device` is not supported for numpy operations') + + x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat) + return 1 - np.clip( (np.dot(x_mat, y_mat.T) + eps) / ( @@ -53,6 +73,9 @@ def sqeuclidean( """ if device is not None: warnings.warn('`device` is not supported for numpy operations') + + x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat) + return ( np.sum(y_mat**2, axis=1) + np.sum(x_mat**2, axis=1)[:, np.newaxis] @@ -78,4 +101,7 @@ def euclidean( """ if device is not None: warnings.warn('`device` is not supported for numpy operations') + + x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat) + return np.sqrt(sqeuclidean(x_mat, y_mat)) From e3d26e6562aa1ff41644194584b44bfd32ee24e3 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Dec 2022 12:09:49 +0100 Subject: [PATCH 08/17] fix: fix metrics and add tests for them Signed-off-by: Johannes Messner --- docarray/utility/find.py | 12 +- .../math/{distance => metrics}/__init__.py | 0 .../math/{distance => metrics}/numpy.py | 34 +++- .../math/{distance => metrics}/torch.py | 25 ++- tests/units/util/__init__.py | 0 tests/units/util/test_find.py | 63 +++++++ tests/units/util/test_metrics.py | 169 ++++++++++++++++++ 7 files changed, 283 insertions(+), 20 deletions(-) rename docarray/utility/math/{distance => metrics}/__init__.py (100%) rename docarray/utility/math/{distance => metrics}/numpy.py (83%) rename docarray/utility/math/{distance => metrics}/torch.py (88%) create mode 100644 tests/units/util/__init__.py create mode 100644 tests/units/util/test_find.py create mode 100644 tests/units/util/test_metrics.py diff --git a/docarray/utility/find.py b/docarray/utility/find.py index ac56ecc16da..eee7c2d048c 100644 --- a/docarray/utility/find.py +++ b/docarray/utility/find.py @@ -18,9 +18,10 @@ def find( index: DocumentArray, query: Union[Tensor, Document, DocumentArray], embedding_field: Optional[str] = 'embedding', - metric: Union[str, Callable[['Tensor', 'Tensor'], 'Tensor']] = 'cosine', + metric: Union[str, Callable[['Tensor', 'Tensor'], 'Tensor']] = 'cosine_sim', limit: int = 10, device: Optional[str] = None, + descending: Optional[bool] = None, ) -> FindResult: """ Find the closest Documents in the index to the query. @@ -48,8 +49,13 @@ def find( :param limit: return the top `limit` results :param device: the computational device to use, can be either `cpu` or a `cuda` device. + :param descending: sort the results in descending order. + Per default, this is chosen based on the `metric` argument. :return: the closest Documents in the index to the query """ + if descending is None: + descending = metric.endswith('_sim') # similarity metrics are descending + embedding_type = _da_attr_type(index, embedding_field) # get framework-specific distance and top_k function @@ -62,7 +68,9 @@ def find( # compute distances and return top results dists = distance_fn(query_embeddings, index_embeddings, device=device) - top_scores, top_indices = top_k_fn(dists, k=limit, device=device) + top_scores, top_indices = top_k_fn( + dists, k=limit, device=device, descending=descending + ) result_docs = [] to_da = True diff --git a/docarray/utility/math/distance/__init__.py b/docarray/utility/math/metrics/__init__.py similarity index 100% rename from docarray/utility/math/distance/__init__.py rename to docarray/utility/math/metrics/__init__.py diff --git a/docarray/utility/math/distance/numpy.py b/docarray/utility/math/metrics/numpy.py similarity index 83% rename from docarray/utility/math/distance/numpy.py rename to docarray/utility/math/metrics/numpy.py index 5aae8a03aaa..78bb607236a 100644 --- a/docarray/utility/math/distance/numpy.py +++ b/docarray/utility/math/metrics/numpy.py @@ -21,7 +21,13 @@ def _expand_if_single_axis(*matrices: 'np.array') -> List['np.array']: return expanded -def cosine( +def _expand_if_scalar(arr: 'np.array') -> 'np.array': + if len(arr.shape) == 0: # avoid scalar output + arr = np.expand_dims(arr, axis=0) + return arr + + +def cosine_sim( x_mat: 'np.ndarray', y_mat: 'np.ndarray', eps: float = 1e-7, @@ -45,18 +51,22 @@ def cosine( x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat) - return 1 - np.clip( + sims = np.clip( (np.dot(x_mat, y_mat.T) + eps) / ( np.outer(np.linalg.norm(x_mat, axis=1), np.linalg.norm(y_mat, axis=1)) + eps ), -1, 1, - ) + ).squeeze() + return _expand_if_scalar(sims) -def sqeuclidean( - x_mat: 'np.ndarray', y_mat: 'np.ndarray', device: Optional[str] = None +def sqeuclidean_dist( + x_mat: 'np.ndarray', + y_mat: 'np.ndarray', + eps: float = 1e-7, + device: Optional[str] = None, ) -> 'np.ndarray': """Pairwise Squared Euclidian distances between all vectors in x_mat and y_mat. @@ -65,6 +75,8 @@ def sqeuclidean( :param y_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is the number of vectors and n_dim is the number of dimensions of each example. :param eps: a small jitter to avoid divde by zero + :param eps: Negative results (due to numerical inaccuracies) no smaller than + -eps will be returned as 0. :param device: Not supported for this backend :return: np.ndarray of shape (n_vectors, n_vectors) containing all pairwise Squared Euclidian distances. @@ -76,14 +88,18 @@ def sqeuclidean( x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat) - return ( + dists = ( np.sum(y_mat**2, axis=1) + np.sum(x_mat**2, axis=1)[:, np.newaxis] - 2 * np.dot(x_mat, y_mat.T) - ) + ).squeeze() + + # remove numerical artifacts + dists = np.where(np.logical_and(dists < 0, dists > -eps), 0, dists) + return _expand_if_scalar(dists) -def euclidean( +def euclidean_dist( x_mat: 'np.ndarray', y_mat: 'np.ndarray', device: Optional[str] = None ) -> 'np.ndarray': """Pairwise Euclidian distances between all vectors in x_mat and y_mat. @@ -104,4 +120,4 @@ def euclidean( x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat) - return np.sqrt(sqeuclidean(x_mat, y_mat)) + return _expand_if_scalar(np.sqrt(sqeuclidean_dist(x_mat, y_mat)).squeeze()) diff --git a/docarray/utility/math/distance/torch.py b/docarray/utility/math/metrics/torch.py similarity index 88% rename from docarray/utility/math/distance/torch.py rename to docarray/utility/math/metrics/torch.py index 4eba79e5441..f0b3942e832 100644 --- a/docarray/utility/math/distance/torch.py +++ b/docarray/utility/math/metrics/torch.py @@ -3,7 +3,6 @@ import torch if TYPE_CHECKING: # pragma: no cover - import numpy from torch import tensor @@ -24,10 +23,16 @@ def _unsqueeze_if_single_axis(*matrices: 'tensor') -> List['tensor']: return unsqueezed -def cosine( +def _usqueeze_if_scalar(t: 'tensor'): + if len(t.shape) == 0: # avoid scalar output + t = t.unsqueeze(0) + return t + + +def cosine_sim( x_mat: 'tensor', y_mat: 'tensor', eps: float = 1e-7, device: Optional[str] = None -) -> 'numpy.ndarray': - """Pairwise cosine distances between all vectors in x_mat and y_mat. +) -> 'tensor': + """Pairwise cosine similarities between all vectors in x_mat and y_mat. :param x_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the number of vectors and n_dim is the number of dimensions of each example. @@ -51,10 +56,11 @@ def cosine( a_n, b_n = x_mat.norm(dim=1)[:, None], y_mat.norm(dim=1)[:, None] a_norm = x_mat / torch.clamp(a_n, min=eps) b_norm = y_mat / torch.clamp(b_n, min=eps) - return 1 - torch.mm(a_norm, b_norm.transpose(0, 1)) + sims = torch.mm(a_norm, b_norm.transpose(0, 1)).squeeze() + return _usqueeze_if_scalar(sims) -def euclidean( +def euclidean_dist( x_mat: 'tensor', y_mat: 'tensor', device: Optional[str] = None ) -> 'tensor': """Pairwise Euclidian distances between all vectors in x_mat and y_mat. @@ -77,10 +83,11 @@ def euclidean( x_mat, y_mat = _unsqueeze_if_single_axis(x_mat, y_mat) - return torch.cdist(x_mat, y_mat) + dists = torch.cdist(x_mat, y_mat).squeeze() + return _usqueeze_if_scalar(dists) -def sqeuclidean( +def sqeuclidean_dist( x_mat: 'tensor', y_mat: 'tensor', device: Optional[str] = None ) -> 'tensor': """Pairwise Squared Euclidian distances between all vectors in x_mat and y_mat. @@ -104,4 +111,4 @@ def sqeuclidean( x_mat, y_mat = _unsqueeze_if_single_axis(x_mat, y_mat) - return torch.cdist(x_mat, y_mat) ** 2 + return _usqueeze_if_scalar((torch.cdist(x_mat, y_mat) ** 2).squeeze()) diff --git a/tests/units/util/__init__.py b/tests/units/util/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/units/util/test_find.py b/tests/units/util/test_find.py new file mode 100644 index 00000000000..4099bf08afe --- /dev/null +++ b/tests/units/util/test_find.py @@ -0,0 +1,63 @@ +import numpy as np +import pytest +import torch + +from docarray import Document, DocumentArray +from docarray.typing import NdArray, TorchTensor +from docarray.utility import find + + +class TorchDoc(Document): + tensor: TorchTensor + + +class NdDoc(Document): + tensor: NdArray + + +@pytest.fixture() +def random_torch_query(): + return TorchDoc(tensor=torch.rand(128)) + + +@pytest.fixture() +def random_nd_query(): + return NdDoc(tensor=np.random.rand(128)) + + +@pytest.fixture() +def random_torch_index(): + return DocumentArray[TorchDoc](TorchDoc(tensor=torch.rand(128)) for _ in range(10)) + + +@pytest.fixture() +def random_nd_index(): + return DocumentArray[NdDoc](NdDoc(tensor=np.random.rand(128)) for _ in range(10)) + + +@pytest.mark.parametrize('metric', ['cosine', 'euclidean', 'sqeuclidean']) +def test_find_torch(random_torch_query, random_torch_index, metric): + top_k, scores = find( + random_torch_index, + random_torch_query, + embedding_field='tensor', + limit=7, + metric=metric, + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert sorted(scores) == scores + + +@pytest.mark.parametrize('metric', ['cosine', 'euclidean', 'sqeuclidean']) +def test_find_np(random_nd_query, random_nd_index, metric): + top_k, scores = find( + random_nd_index, + random_nd_query, + embedding_field='tensor', + limit=7, + metric=metric, + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert sorted(scores) == scores diff --git a/tests/units/util/test_metrics.py b/tests/units/util/test_metrics.py new file mode 100644 index 00000000000..c50af26fc76 --- /dev/null +++ b/tests/units/util/test_metrics.py @@ -0,0 +1,169 @@ +import numpy as np +import torch + +import docarray.utility.math.metrics.numpy as np_distances +import docarray.utility.math.metrics.torch as torch_distances + + +def test_cosine_sim_torch(): + a = torch.rand(128) + b = torch.rand(128) + assert torch_distances.cosine_sim(a, b).shape == (1,) + assert torch_distances.cosine_sim(a, b) == torch_distances.cosine_sim(b, a) + torch.testing.assert_close(torch_distances.cosine_sim(a, a), torch.ones(1)) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + assert torch_distances.cosine_sim(a, b).shape == (10, 5) + assert torch_distances.cosine_sim(b, a).shape == (5, 10) + diag_dists = torch.diagonal(torch_distances.cosine_sim(b, b)) # self-comparisons + torch.testing.assert_allclose(diag_dists, torch.ones(5)) + + +def test_cosine_sim_np(): + a = np.random.rand(128) + b = np.random.rand(128) + assert np_distances.cosine_sim(a, b).shape == (1,) + assert np_distances.cosine_sim(a, b) == np_distances.cosine_sim(b, a) + np.testing.assert_array_almost_equal(np_distances.cosine_sim(a, a), np.ones((1,))) + + a = np.random.rand(10, 3) + b = np.random.rand(5, 3) + assert np_distances.cosine_sim(a, b).shape == (10, 5) + assert np_distances.cosine_sim(b, a).shape == (5, 10) + diag_dists = np.diagonal(np_distances.cosine_sim(b, b)) # self-comparisons + np.testing.assert_array_almost_equal(diag_dists, np.ones((5,))) + + +def test_cosine_sim_compare(): + a = torch.rand(128) + b = torch.rand(128) + torch.testing.assert_close( + torch_distances.cosine_sim(a, b), + torch.from_numpy(np_distances.cosine_sim(a.numpy(), b.numpy())), + ) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + torch.testing.assert_close( + torch_distances.cosine_sim(a, b), + torch.from_numpy(np_distances.cosine_sim(a.numpy(), b.numpy())), + ) + + +def test_euclidean_dist_torch(): + a = torch.rand(128) + b = torch.rand(128) + assert torch_distances.euclidean_dist(a, b).shape == (1,) + assert torch_distances.euclidean_dist(a, b) == torch_distances.euclidean_dist(b, a) + torch.testing.assert_close(torch_distances.euclidean_dist(a, a), torch.zeros(1)) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + assert torch_distances.euclidean_dist(a, b).shape == (10, 5) + assert torch_distances.euclidean_dist(b, a).shape == (5, 10) + diag_dists = torch.diagonal( + torch_distances.euclidean_dist(b, b) + ) # self-comparisons + torch.testing.assert_allclose(diag_dists, torch.zeros(5)) + + a = torch.tensor([0.0, 2.0, 0.0]) + b = torch.tensor([0.0, 0.0, 2.0]) + desired_output_singleton = torch.sqrt(torch.tensor([2.0**2.0 + 2.0**2.0])) + torch.testing.assert_close( + torch_distances.euclidean_dist(a, b), desired_output_singleton + ) + + a = torch.tensor([[0.0, 2.0, 0.0], [0.0, 0.0, 2.0]]) + b = torch.tensor([[0.0, 0.0, 2.0], [0.0, 2.0, 0.0]]) + desired_output_singleton = torch.tensor( + [[desired_output_singleton.item(), 0.0], [0.0, desired_output_singleton.item()]] + ) + torch.testing.assert_close( + torch_distances.euclidean_dist(a, b), desired_output_singleton + ) + + +def test_euclidean_dist_np(): + a = np.random.rand(128) + b = np.random.rand(128) + assert np_distances.euclidean_dist(a, b).shape == (1,) + assert np_distances.euclidean_dist(a, b) == np_distances.euclidean_dist(b, a) + np.testing.assert_array_almost_equal( + np_distances.euclidean_dist(a, a), np.zeros((1,)) + ) + + a = np.random.rand(10, 3) + b = np.random.rand(5, 3) + assert np_distances.euclidean_dist(a, b).shape == (10, 5) + assert np_distances.euclidean_dist(b, a).shape == (5, 10) + diag_dists = np.diagonal(np_distances.euclidean_dist(b, b)) # self-comparisons + np.testing.assert_array_almost_equal(diag_dists, np.zeros((5,))) + + a = np.array([0.0, 2.0, 0.0]) + b = np.array([0.0, 0.0, 2.0]) + desired_output_singleton = np.sqrt(np.array([2.0**2.0 + 2.0**2.0])) + np.testing.assert_array_almost_equal( + np_distances.euclidean_dist(a, b), desired_output_singleton + ) + + a = np.array([[0.0, 2.0, 0.0], [0.0, 0.0, 2.0]]) + b = np.array([[0.0, 0.0, 2.0], [0.0, 2.0, 0.0]]) + desired_output_singleton = np.array( + [[desired_output_singleton.item(), 0.0], [0.0, desired_output_singleton.item()]] + ) + np.testing.assert_array_almost_equal( + np_distances.euclidean_dist(a, b), desired_output_singleton + ) + + +def test_euclidean_dist_compare(): + a = torch.rand(128) + b = torch.rand(128) + torch.testing.assert_close( + torch_distances.euclidean_dist(a, b), + torch.from_numpy(np_distances.euclidean_dist(a.numpy(), b.numpy())).to( + torch.float32 + ), + ) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + torch.testing.assert_close( + torch_distances.euclidean_dist(a, b), + torch.from_numpy(np_distances.euclidean_dist(a.numpy(), b.numpy())), + ) + + +def test_sqeuclidean_dist_torch(): + a = torch.rand(128) + b = torch.rand(128) + assert torch_distances.sqeuclidean_dist(a, b).shape == (1,) + torch.testing.assert_close( + torch_distances.sqeuclidean_dist(a, b), + torch_distances.euclidean_dist(a, b) ** 2, + ) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + assert torch_distances.sqeuclidean_dist(a, b).shape == (10, 5) + torch.testing.assert_close( + torch_distances.sqeuclidean_dist(a, b), + torch_distances.euclidean_dist(a, b) ** 2, + ) + + +def test_sqeuclidea_dist_np(): + a = np.random.rand(128) + b = np.random.rand(128) + assert np_distances.sqeuclidean_dist(a, b).shape == (1,) + np.testing.assert_array_almost_equal( + np_distances.sqeuclidean_dist(a, b), np_distances.euclidean_dist(a, b) ** 2 + ) + + a = np.random.rand(10, 3) + b = np.random.rand(5, 3) + assert np_distances.sqeuclidean_dist(a, b).shape == (10, 5) + np.testing.assert_array_almost_equal( + np_distances.sqeuclidean_dist(a, b), np_distances.euclidean_dist(a, b) ** 2 + ) From 1ed7d93029fe0fad05a50b86fdb81a80756f3132 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Dec 2022 14:45:13 +0100 Subject: [PATCH 09/17] test: add tests for find Signed-off-by: Johannes Messner --- docarray/utility/find.py | 6 +++--- docarray/utility/helper/numpy.py | 5 ++++- docarray/utility/helper/torch.py | 3 ++- tests/units/util/test_find.py | 14 ++++++++++---- 4 files changed, 19 insertions(+), 9 deletions(-) diff --git a/docarray/utility/find.py b/docarray/utility/find.py index eee7c2d048c..8add92b9c17 100644 --- a/docarray/utility/find.py +++ b/docarray/utility/find.py @@ -75,9 +75,9 @@ def find( result_docs = [] to_da = True for top_idx in top_indices: # workaround until #930 is fixed - if len(top_idx) == 0: + if top_idx.shape == () or len(top_idx) == 0: # single result for single query result_docs.append(index[top_idx]) - else: + else: # there were multiple queries, so multiple results inner_result_docs = [] to_da = False for inner_top_idx in top_idx: @@ -173,6 +173,6 @@ def _get_distance_fn(embedding_type: Type, distance_name: str) -> Callable: """ framework = type_to_framework[embedding_type] return getattr( - importlib.import_module(f'docarray.utility.math.distance.{framework}'), + importlib.import_module(f'docarray.utility.math.metrics.{framework}'), f'{distance_name}', ) diff --git a/docarray/utility/helper/numpy.py b/docarray/utility/helper/numpy.py index d4d41f72a03..99cbf8c05c7 100644 --- a/docarray/utility/helper/numpy.py +++ b/docarray/utility/helper/numpy.py @@ -26,6 +26,9 @@ def top_k( if device is not None: warnings.warn('`device` is not supported for numpy operations') + if len(values.shape) == 1: + values = np.expand_dims(values, axis=0) + if descending: values = -values @@ -42,4 +45,4 @@ def top_k( if descending: values = -values - return values, idx + return values.squeeze(), idx.squeeze() diff --git a/docarray/utility/helper/torch.py b/docarray/utility/helper/torch.py index f23a19c58a3..e0e8b89174e 100644 --- a/docarray/utility/helper/torch.py +++ b/docarray/utility/helper/torch.py @@ -24,5 +24,6 @@ def top_k( """ if device is not None: values = values.to(device) - k = max(k, len(values[-1])) + len_values = values.shape[-1] if len(values.shape) > 1 else len(values) + k = min(k, len_values) return torch.topk(input=values, k=k, largest=descending, sorted=True, dim=-1) diff --git a/tests/units/util/test_find.py b/tests/units/util/test_find.py index 4099bf08afe..7aa6ce5f856 100644 --- a/tests/units/util/test_find.py +++ b/tests/units/util/test_find.py @@ -35,7 +35,7 @@ def random_nd_index(): return DocumentArray[NdDoc](NdDoc(tensor=np.random.rand(128)) for _ in range(10)) -@pytest.mark.parametrize('metric', ['cosine', 'euclidean', 'sqeuclidean']) +@pytest.mark.parametrize('metric', ['cosine_sim', 'euclidean_dist', 'sqeuclidean_dist']) def test_find_torch(random_torch_query, random_torch_index, metric): top_k, scores = find( random_torch_index, @@ -46,10 +46,13 @@ def test_find_torch(random_torch_query, random_torch_index, metric): ) assert len(top_k) == 7 assert len(scores) == 7 - assert sorted(scores) == scores + if metric.endswith('_dist'): + assert (torch.stack(sorted(scores)) == scores).all() + else: + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() -@pytest.mark.parametrize('metric', ['cosine', 'euclidean', 'sqeuclidean']) +@pytest.mark.parametrize('metric', ['cosine_sim', 'euclidean_dist', 'sqeuclidean_dist']) def test_find_np(random_nd_query, random_nd_index, metric): top_k, scores = find( random_nd_index, @@ -60,4 +63,7 @@ def test_find_np(random_nd_query, random_nd_index, metric): ) assert len(top_k) == 7 assert len(scores) == 7 - assert sorted(scores) == scores + if metric.endswith('_dist'): + assert (sorted(scores) == scores).all() + else: + assert (sorted(scores, reverse=True) == scores).all() From 0f7068ff1f8bab1a27fd5d95d8824f93bc3afecf Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Dec 2022 15:19:09 +0100 Subject: [PATCH 10/17] test: add test for topk Signed-off-by: Johannes Messner --- docarray/utility/find.py | 56 ++++++++++++++++++++++++++++++-- docarray/utility/helper/numpy.py | 10 +++--- docarray/utility/helper/torch.py | 9 +++-- tests/units/util/test_topk.py | 49 ++++++++++++++++++++++++++++ 4 files changed, 115 insertions(+), 9 deletions(-) create mode 100644 tests/units/util/test_topk.py diff --git a/docarray/utility/find.py b/docarray/utility/find.py index 8add92b9c17..ebf56e369ad 100644 --- a/docarray/utility/find.py +++ b/docarray/utility/find.py @@ -9,20 +9,40 @@ class FindResult(NamedTuple): - # for single query | for multiple queries - documents: Union[DocumentArray, List[DocumentArray]] + documents: DocumentArray scores: Tensor def find( index: DocumentArray, - query: Union[Tensor, Document, DocumentArray], + query: Union[Tensor, Document], embedding_field: Optional[str] = 'embedding', metric: Union[str, Callable[['Tensor', 'Tensor'], 'Tensor']] = 'cosine_sim', limit: int = 10, device: Optional[str] = None, descending: Optional[bool] = None, ) -> FindResult: + query = _extract_embedding_single(query, embedding_field) + return find_batched( + index=index, + query=query, + embedding_field=embedding_field, + metric=metric, + limit=limit, + device=device, + descending=descending, + )[0] + + +def find_batched( + index: DocumentArray, + query: Union[Tensor, DocumentArray], + embedding_field: Optional[str] = 'embedding', + metric: Union[str, Callable[['Tensor', 'Tensor'], 'Tensor']] = 'cosine_sim', + limit: int = 10, + device: Optional[str] = None, + descending: Optional[bool] = None, +) -> List[FindResult]: """ Find the closest Documents in the index to the query. @@ -72,6 +92,7 @@ def find( dists, k=limit, device=device, descending=descending ) + # results = [] result_docs = [] to_da = True for top_idx in top_indices: # workaround until #930 is fixed @@ -88,6 +109,35 @@ def find( return FindResult(documents=result_docs, scores=top_scores) +def _extract_embedding_single( + data: Union[DocumentArray, Document, Tensor], + embedding_field: str, +) -> Tensor: + """Extract the embeddings from a single query, + and return it in a batched representation. + + :param data: the data + :param embedding_field: the embedding field + :param embedding_type: type of the embedding: torch.Tensor, numpy.ndarray etc. + :return: the embeddings + """ + if isinstance(data, Document): + emb = getattr(data, embedding_field) + else: # treat data as tensor + emb = data + if len(emb.shape) == 1: + # TODO(johannes) solve this with computational backend, + # this is ugly hack for now + if isinstance(emb, torch.Tensor): + emb = emb.unsqueeze(0) + else: + import numpy as np + + if isinstance(emb, np.ndarray): + emb = np.expand_dims(emb, axis=0) + return emb + + def _extraxt_embeddings( data: Union[DocumentArray, Document, Tensor], embedding_field: str, diff --git a/docarray/utility/helper/numpy.py b/docarray/utility/helper/numpy.py index 99cbf8c05c7..a3350bcfd69 100644 --- a/docarray/utility/helper/numpy.py +++ b/docarray/utility/helper/numpy.py @@ -16,12 +16,14 @@ def top_k( by setting the `descending` flag. :param values: Torch tensor of values to rank. - Should be a vector (single axis), or squeezable. - TODO(johannes) is squeezable a term understood in np land? + Should be of shape (n_queries, n_values_per_query). + Inputs of shape (n_values_per_query,) will be expanded + to (1, n_values_per_query). :param k: number of values to retrieve :param descending: retrieve largest values instead of smallest values :param device: Not supported for this backend - :return: Tuple containing the retrieved values, and their indices + :return: Tuple containing the retrieved values, and their indices. + Both ar of shape (n_queries, k) """ if device is not None: warnings.warn('`device` is not supported for numpy operations') @@ -45,4 +47,4 @@ def top_k( if descending: values = -values - return values.squeeze(), idx.squeeze() + return values, idx diff --git a/docarray/utility/helper/torch.py b/docarray/utility/helper/torch.py index e0e8b89174e..0ceeee20101 100644 --- a/docarray/utility/helper/torch.py +++ b/docarray/utility/helper/torch.py @@ -15,15 +15,20 @@ def top_k( by setting the `descending` flag. :param values: Torch tensor of values to rank. - Has to be of shape (n_queries, n_values_per_query) + Should be of shape (n_queries, n_values_per_query). + Inputs of shape (n_values_per_query,) will be expanded + to (1, n_values_per_query). :param k: number of values to retrieve :param descending: retrieve largest values instead of smallest values :param device: the computational device to use, can be either `cpu` or a `cuda` device. - :return: Tuple containing the retrieved values, and their indices + :return: Tuple containing the retrieved values, and their indices. + Both ar of shape (n_queries, k) """ if device is not None: values = values.to(device) + if len(values.shape) <= 1: + values = values.view(1, -1) len_values = values.shape[-1] if len(values.shape) > 1 else len(values) k = min(k, len_values) return torch.topk(input=values, k=k, largest=descending, sorted=True, dim=-1) diff --git a/tests/units/util/test_topk.py b/tests/units/util/test_topk.py new file mode 100644 index 00000000000..2bd8d3b520a --- /dev/null +++ b/tests/units/util/test_topk.py @@ -0,0 +1,49 @@ +import numpy as np +import torch + +from docarray.utility.helper.numpy import top_k as np_top_k +from docarray.utility.helper.torch import top_k as torch_top_k + + +def test_topk_torch(): + a = torch.tensor([1, 4, 2, 7, 4, 9, 2]) + vals, indices = torch_top_k(a, 3) + assert vals.shape == (1, 3) + assert indices.shape == (1, 3) + assert (vals.squeeze() == torch.tensor([1, 2, 2])).all() + assert (indices.squeeze() == torch.tensor([0, 2, 6])).all() or ( + indices.squeeze() == torch.tensor([0, 6, 2]) + ).all() + + a = torch.tensor([[1, 4, 2, 7, 4, 9, 2], [11, 6, 2, 7, 3, 10, 4]]) + vals, indices = torch_top_k(a, 3) + assert vals.shape == (2, 3) + assert indices.shape == (2, 3) + assert (vals[0] == torch.tensor([1, 2, 2])).all() + assert (indices[0] == torch.tensor([0, 2, 6])).all() or ( + indices[0] == torch.tensor([0, 6, 2]) + ).all() + assert (vals[1] == torch.tensor([2, 3, 4])).all() + assert (indices[1] == torch.tensor([2, 4, 6])).all() + + +def test_topk_numpy(): + a = np.array([1, 4, 2, 7, 4, 9, 2]) + vals, indices = np_top_k(a, 3) + assert vals.shape == (1, 3) + assert indices.shape == (1, 3) + assert (vals.squeeze() == np.array([1, 2, 2])).all() + assert (indices.squeeze() == np.array([0, 2, 6])).all() or ( + indices.squeeze() == np.array([0, 6, 2]) + ).all() + + a = np.array([[1, 4, 2, 7, 4, 9, 2], [11, 6, 2, 7, 3, 10, 4]]) + vals, indices = np_top_k(a, 3) + assert vals.shape == (2, 3) + assert indices.shape == (2, 3) + assert (vals[0] == np.array([1, 2, 2])).all() + assert (indices[0] == np.array([0, 2, 6])).all() or ( + indices[0] == np.array([0, 6, 2]) + ).all() + assert (vals[1] == np.array([2, 3, 4])).all() + assert (indices[1] == np.array([2, 4, 6])).all() From d5c9ef9701c89115158ed337381bf672b18ed22f Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Dec 2022 16:03:08 +0100 Subject: [PATCH 11/17] feat: add batched find Signed-off-by: Johannes Messner --- docarray/utility/__init__.py | 4 +-- docarray/utility/find.py | 39 ++++++-------------------- tests/units/util/test_find.py | 52 ++++++++++++++++++++++++++++++++++- 3 files changed, 61 insertions(+), 34 deletions(-) diff --git a/docarray/utility/__init__.py b/docarray/utility/__init__.py index 69148798f93..7c9bada7262 100644 --- a/docarray/utility/__init__.py +++ b/docarray/utility/__init__.py @@ -1,3 +1,3 @@ -from docarray.utility.find import find +from docarray.utility.find import find, find_batched -__all__ = ['find'] +__all__ = ['find', 'find_batched'] diff --git a/docarray/utility/find.py b/docarray/utility/find.py index ebf56e369ad..a72d52767a2 100644 --- a/docarray/utility/find.py +++ b/docarray/utility/find.py @@ -92,21 +92,14 @@ def find_batched( dists, k=limit, device=device, descending=descending ) - # results = [] - result_docs = [] - to_da = True - for top_idx in top_indices: # workaround until #930 is fixed - if top_idx.shape == () or len(top_idx) == 0: # single result for single query - result_docs.append(index[top_idx]) - else: # there were multiple queries, so multiple results - inner_result_docs = [] - to_da = False - for inner_top_idx in top_idx: - inner_result_docs.append(index[inner_top_idx]) - if to_da: - result_docs = DocumentArray(result_docs) - - return FindResult(documents=result_docs, scores=top_scores) + results = [] + for indices_per_query, scores_per_query in zip(top_indices, top_scores): + docs_per_query = [] + for idx in indices_per_query: # workaround until #930 is fixed + docs_per_query.append(index[idx]) + docs_per_query = DocumentArray(docs_per_query) + results.append(FindResult(scores=scores_per_query, documents=docs_per_query)) + return results def _extract_embedding_single( @@ -173,22 +166,6 @@ def _extraxt_embeddings( return emb -def _to_documentarray_query( - query: Union[Tensor, Document, DocumentArray], embedding_field: str -) -> DocumentArray: - """Convert the query to a DocumentArray. - - :param query: the query - :return: the query as DocumentArray - """ - if isinstance(query, DocumentArray): - return query - elif isinstance(query, Document): - return DocumentArray([query]) - else: - return DocumentArray([Document(**{embedding_field: query})]) - - def _da_attr_type(da: DocumentArray, attr: str) -> Type: """Get the type of the attribute according to the Document type (schema) of the DocumentArray. diff --git a/tests/units/util/test_find.py b/tests/units/util/test_find.py index 7aa6ce5f856..ef4382bd1e6 100644 --- a/tests/units/util/test_find.py +++ b/tests/units/util/test_find.py @@ -4,7 +4,7 @@ from docarray import Document, DocumentArray from docarray.typing import NdArray, TorchTensor -from docarray.utility import find +from docarray.utility import find, find_batched class TorchDoc(Document): @@ -20,11 +20,21 @@ def random_torch_query(): return TorchDoc(tensor=torch.rand(128)) +@pytest.fixture() +def random_torch_batch_query(): + return DocumentArray[TorchDoc]([TorchDoc(tensor=torch.rand(128)) for _ in range(5)]) + + @pytest.fixture() def random_nd_query(): return NdDoc(tensor=np.random.rand(128)) +@pytest.fixture() +def random_nd_batch_query(): + return DocumentArray[NdDoc]([NdDoc(tensor=np.random.rand(128)) for _ in range(5)]) + + @pytest.fixture() def random_torch_index(): return DocumentArray[TorchDoc](TorchDoc(tensor=torch.rand(128)) for _ in range(10)) @@ -67,3 +77,43 @@ def test_find_np(random_nd_query, random_nd_index, metric): assert (sorted(scores) == scores).all() else: assert (sorted(scores, reverse=True) == scores).all() + + +@pytest.mark.parametrize('metric', ['cosine_sim', 'euclidean_dist', 'sqeuclidean_dist']) +def test_find_batched_torch(random_torch_batch_query, random_torch_index, metric): + results = find_batched( + random_torch_index, + random_torch_batch_query, + embedding_field='tensor', + limit=7, + metric=metric, + ) + assert len(results) == len(random_torch_batch_query) + for top_k, scores in results: + assert len(top_k) == 7 + assert len(scores) == 7 + for sc in [scores for _, scores in results]: + if metric.endswith('_dist'): + assert (torch.stack(sorted(sc)) == sc).all() + else: + assert (torch.stack(sorted(sc, reverse=True)) == sc).all() + + +@pytest.mark.parametrize('metric', ['cosine_sim', 'euclidean_dist', 'sqeuclidean_dist']) +def test_find_batched_np(random_nd_batch_query, random_nd_index, metric): + results = find_batched( + random_nd_index, + random_nd_batch_query, + embedding_field='tensor', + limit=7, + metric=metric, + ) + assert len(results) == len(random_nd_batch_query) + for top_k, scores in results: + assert len(top_k) == 7 + assert len(scores) == 7 + for sc in [scores for _, scores in results]: + if metric.endswith('_dist'): + assert (sorted(sc) == sc).all() + else: + assert (sorted(sc, reverse=True) == sc).all() From 11cb7f28b42f13344bcb8247e7b2d79777ab9095 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Dec 2022 16:30:52 +0100 Subject: [PATCH 12/17] test: add more tests Signed-off-by: Johannes Messner --- tests/units/util/test_find.py | 135 ++++++++++++++++++++++++++++++++++ 1 file changed, 135 insertions(+) diff --git a/tests/units/util/test_find.py b/tests/units/util/test_find.py index ef4382bd1e6..665b880912c 100644 --- a/tests/units/util/test_find.py +++ b/tests/units/util/test_find.py @@ -62,6 +62,34 @@ def test_find_torch(random_torch_query, random_torch_index, metric): assert (torch.stack(sorted(scores, reverse=True)) == scores).all() +def test_find_torch_tensor_query(random_torch_query, random_torch_index): + query = random_torch_query.tensor + top_k, scores = find( + random_torch_index, + query, + embedding_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() + + +def test_find_torch_stacked(random_torch_query, random_torch_index): + random_torch_index.stack() + top_k, scores = find( + random_torch_index, + random_torch_query, + embedding_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() + + @pytest.mark.parametrize('metric', ['cosine_sim', 'euclidean_dist', 'sqeuclidean_dist']) def test_find_np(random_nd_query, random_nd_index, metric): top_k, scores = find( @@ -79,6 +107,34 @@ def test_find_np(random_nd_query, random_nd_index, metric): assert (sorted(scores, reverse=True) == scores).all() +def test_find_np_tensor_query(random_nd_query, random_nd_index): + query = random_nd_query.tensor + top_k, scores = find( + random_nd_index, + query, + embedding_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (sorted(scores, reverse=True) == scores).all() + + +def test_find_np_stacked(random_nd_query, random_nd_index): + random_nd_index.stack() + top_k, scores = find( + random_nd_index, + random_nd_query, + embedding_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (sorted(scores, reverse=True) == scores).all() + + @pytest.mark.parametrize('metric', ['cosine_sim', 'euclidean_dist', 'sqeuclidean_dist']) def test_find_batched_torch(random_torch_batch_query, random_torch_index, metric): results = find_batched( @@ -99,6 +155,47 @@ def test_find_batched_torch(random_torch_batch_query, random_torch_index, metric assert (torch.stack(sorted(sc, reverse=True)) == sc).all() +def test_find_batched_torch_tensor_query(random_torch_batch_query, random_torch_index): + query = torch.stack(random_torch_batch_query.tensor) + results = find_batched( + random_torch_index, + query, + embedding_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(results) == len(random_torch_batch_query) + for top_k, scores in results: + assert len(top_k) == 7 + assert len(scores) == 7 + for sc in [scores for _, scores in results]: + assert (torch.stack(sorted(sc, reverse=True)) == sc).all() + + +@pytest.mark.parametrize('stack_what', ['index', 'query', 'both']) +def test_find_batched_torch_stacked( + random_torch_batch_query, random_torch_index, stack_what +): + if stack_what in ('index', 'both'): + random_torch_index.stack() + if stack_what in ('query', 'both'): + random_torch_batch_query.stack() + + results = find_batched( + random_torch_index, + random_torch_batch_query, + embedding_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(results) == len(random_torch_batch_query) + for top_k, scores in results: + assert len(top_k) == 7 + assert len(scores) == 7 + for sc in [scores for _, scores in results]: + assert (torch.stack(sorted(sc, reverse=True)) == sc).all() + + @pytest.mark.parametrize('metric', ['cosine_sim', 'euclidean_dist', 'sqeuclidean_dist']) def test_find_batched_np(random_nd_batch_query, random_nd_index, metric): results = find_batched( @@ -117,3 +214,41 @@ def test_find_batched_np(random_nd_batch_query, random_nd_index, metric): assert (sorted(sc) == sc).all() else: assert (sorted(sc, reverse=True) == sc).all() + + +def test_find_batched_np_tensor_query(random_nd_batch_query, random_nd_index): + query = np.stack(random_nd_batch_query.tensor) + results = find_batched( + random_nd_index, + query, + embedding_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(results) == len(random_nd_batch_query) + for top_k, scores in results: + assert len(top_k) == 7 + assert len(scores) == 7 + for sc in [scores for _, scores in results]: + assert (sorted(sc, reverse=True) == sc).all() + + +@pytest.mark.parametrize('stack_what', ['index', 'query', 'both']) +def test_find_batched_np_stacked(random_nd_batch_query, random_nd_index, stack_what): + if stack_what in ('index', 'both'): + random_nd_index.stack() + if stack_what in ('query', 'both'): + random_nd_batch_query.stack() + results = find_batched( + random_nd_index, + random_nd_batch_query, + embedding_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(results) == len(random_nd_batch_query) + for top_k, scores in results: + assert len(top_k) == 7 + assert len(scores) == 7 + for sc in [scores for _, scores in results]: + assert (sorted(sc, reverse=True) == sc).all() From ba56987067971624df7e3e9fce7fcbf216c2e0c0 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Dec 2022 16:35:35 +0100 Subject: [PATCH 13/17] docs: improve docstrings Signed-off-by: Johannes Messner --- docarray/utility/find.py | 38 ++++++++++++++++++++++++-- docarray/utility/helper/numpy.py | 3 +- docarray/utility/helper/torch.py | 3 +- docarray/utility/math/metrics/numpy.py | 2 +- 4 files changed, 41 insertions(+), 5 deletions(-) diff --git a/docarray/utility/find.py b/docarray/utility/find.py index a72d52767a2..32813bb540e 100644 --- a/docarray/utility/find.py +++ b/docarray/utility/find.py @@ -22,6 +22,38 @@ def find( device: Optional[str] = None, descending: Optional[bool] = None, ) -> FindResult: + """ + Find the closest Documents in the index to the query. + + .. note:: + This utility function is likely to be removed once + Document Stores are available. + At that point, and in-memory Document Store will serve the same purpose + by exposing a .find() method. + + .. note:: + This is a simple implementation that assumes the same embedding field name for + both query and index, does not support nested search, and does not support + hybrid (multi-vector) search. These shortcoming will be addressed in future + versions. + + :param index: the index of Documents to search in + :param query: the query to search for + :param embedding_field: the tensor-like field in the index to use + for the similarity computation + :param metric: the distance metric to use for the similarity computation. + Can be a string specifying a predefined metric + (TODO(johannes) specify) or a callable that takes two tensors and returns + a tensor of distances. TODO(johannes) implement passing a callable + :param limit: return the top `limit` results + :param device: the computational device to use, + can be either `cpu` or a `cuda` device. + :param descending: sort the results in descending order. + Per default, this is chosen based on the `metric` argument. + :return: A named tuple of the form (DocumentArray, Tensor), + where the first element contains the closes matches for the query, + and the second element contains the corresponding scores. + """ query = _extract_embedding_single(query, embedding_field) return find_batched( index=index, @@ -44,7 +76,7 @@ def find_batched( descending: Optional[bool] = None, ) -> List[FindResult]: """ - Find the closest Documents in the index to the query. + Find the closest Documents in the index to the queries. .. note:: This utility function is likely to be removed once @@ -71,7 +103,9 @@ def find_batched( can be either `cpu` or a `cuda` device. :param descending: sort the results in descending order. Per default, this is chosen based on the `metric` argument. - :return: the closest Documents in the index to the query + :return: a list of named tuples of the form (DocumentArray, Tensor), + where the first element contains the closes matches for each query, + and the second element contains the corresponding scores. """ if descending is None: descending = metric.endswith('_sim') # similarity metrics are descending diff --git a/docarray/utility/helper/numpy.py b/docarray/utility/helper/numpy.py index a3350bcfd69..cf2dcadc222 100644 --- a/docarray/utility/helper/numpy.py +++ b/docarray/utility/helper/numpy.py @@ -10,7 +10,8 @@ def top_k( descending: bool = False, device: Optional[str] = None, ) -> Tuple['np.ndarray', 'np.ndarray']: - """Retrieves the top k smallest values in `values`, + """ + Retrieves the top k smallest values in `values`, and returns them alongside their indices in the input `values`. Can also be used to retrieve the top k largest values, by setting the `descending` flag. diff --git a/docarray/utility/helper/torch.py b/docarray/utility/helper/torch.py index 0ceeee20101..adf1fe8f28c 100644 --- a/docarray/utility/helper/torch.py +++ b/docarray/utility/helper/torch.py @@ -9,7 +9,8 @@ def top_k( descending: bool = False, device: Optional[str] = None, ) -> Tuple['torch.tensor', 'torch.tensor']: - """Retrieves the top k smallest values in `values`, + """ + Retrieves the top k smallest values in `values`, and returns them alongside their indices in the input `values`. Can also be used to retrieve the top k largest values, by setting the `descending` flag. diff --git a/docarray/utility/math/metrics/numpy.py b/docarray/utility/math/metrics/numpy.py index 78bb607236a..b9e1810b4d4 100644 --- a/docarray/utility/math/metrics/numpy.py +++ b/docarray/utility/math/metrics/numpy.py @@ -33,7 +33,7 @@ def cosine_sim( eps: float = 1e-7, device: Optional[str] = None, ) -> 'np.ndarray': - """Pairwise cosine distances between all vectors in x_mat and y_mat. + """Pairwise cosine similarities between all vectors in x_mat and y_mat. :param x_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is the number of vectors and n_dim is the number of dimensions of each example. From 840ef794c25484f83d158996bd4befe630713dae Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Dec 2022 17:24:18 +0100 Subject: [PATCH 14/17] docs: improve docstrings Signed-off-by: Johannes Messner --- docarray/utility/find.py | 108 +++++++++++++++++++++++++++++++++------ 1 file changed, 93 insertions(+), 15 deletions(-) diff --git a/docarray/utility/find.py b/docarray/utility/find.py index 32813bb540e..32d58553bee 100644 --- a/docarray/utility/find.py +++ b/docarray/utility/find.py @@ -16,14 +16,15 @@ class FindResult(NamedTuple): def find( index: DocumentArray, query: Union[Tensor, Document], - embedding_field: Optional[str] = 'embedding', - metric: Union[str, Callable[['Tensor', 'Tensor'], 'Tensor']] = 'cosine_sim', + embedding_field: str = 'embedding', + metric: str = 'cosine_sim', limit: int = 10, device: Optional[str] = None, descending: Optional[bool] = None, ) -> FindResult: """ Find the closest Documents in the index to the query. + Supports PyTorch and NumPy embeddings. .. note:: This utility function is likely to be removed once @@ -37,14 +38,49 @@ def find( hybrid (multi-vector) search. These shortcoming will be addressed in future versions. + EXAMPLE USAGE + + .. code-block:: python + + from docarray import DocumentArray, Document + from docarray.typing import TorchTensor + from docarray.utility.find import find + + + class MyDocument(Document): + embedding: TorchTensor + + + index = DocumentArray[MyDocument]( + [MyDocument(embedding=torch.rand(128)) for _ in range(100)] + ) + + # use Document as query + query = MyDocument(embedding=torch.rand(128)) + top_matches, scores = find( + index=index, + query=query, + embedding_field='tensor', + metric='cosine_sim', + ) + + # use tensor as query + query = torch.rand(128) + top_matches, scores = find( + index=index, + query=query, + embedding_field='tensor', + metric='cosine_sim', + ) + :param index: the index of Documents to search in :param query: the query to search for :param embedding_field: the tensor-like field in the index to use for the similarity computation :param metric: the distance metric to use for the similarity computation. - Can be a string specifying a predefined metric - (TODO(johannes) specify) or a callable that takes two tensors and returns - a tensor of distances. TODO(johannes) implement passing a callable + Can be one of the following strings: + 'cosine_sim' for cosine similarity, 'euclidean_dist' for euclidean distance, + 'sqeuclidean_dist' for squared euclidean distance :param limit: return the top `limit` results :param device: the computational device to use, can be either `cpu` or a `cuda` device. @@ -69,7 +105,7 @@ def find( def find_batched( index: DocumentArray, query: Union[Tensor, DocumentArray], - embedding_field: Optional[str] = 'embedding', + embedding_field: str = 'embedding', metric: Union[str, Callable[['Tensor', 'Tensor'], 'Tensor']] = 'cosine_sim', limit: int = 10, device: Optional[str] = None, @@ -77,6 +113,7 @@ def find_batched( ) -> List[FindResult]: """ Find the closest Documents in the index to the queries. + Supports PyTorch and NumPy embeddings. .. note:: This utility function is likely to be removed once @@ -90,14 +127,53 @@ def find_batched( hybrid (multi-vector) search. These shortcoming will be addressed in future versions. + EXAMPLE USAGE + + .. code-block:: python + + from docarray import DocumentArray, Document + from docarray.typing import TorchTensor + from docarray.utility.find import find + + + class MyDocument(Document): + embedding: TorchTensor + + + index = DocumentArray[MyDocument]( + [MyDocument(embedding=torch.rand(128)) for _ in range(100)] + ) + + # use DocumentArray as query + query = DocumentArray[MyDocument]( + [MyDocument(embedding=torch.rand(128)) for _ in range(3)] + ) + results = find( + index=index, + query=query, + embedding_field='tensor', + metric='cosine_sim', + ) + top_matches, scores = results[0] + + # use tensor as query + query = torch.rand(3, 128) + results, scores = find( + index=index, + query=query, + embedding_field='tensor', + metric='cosine_sim', + ) + top_matches, scores = results[0] + :param index: the index of Documents to search in :param query: the query to search for :param embedding_field: the tensor-like field in the index to use for the similarity computation :param metric: the distance metric to use for the similarity computation. - Can be a string specifying a predefined metric - (TODO(johannes) specify) or a callable that takes two tensors and returns - a tensor of distances. TODO(johannes) implement passing a callable + Can be one of the following strings: + 'cosine_sim' for cosine similarity, 'euclidean_dist' for euclidean distance, + 'sqeuclidean_dist' for squared euclidean distance :param limit: return the top `limit` results :param device: the computational device to use, can be either `cpu` or a `cuda` device. @@ -113,7 +189,7 @@ def find_batched( embedding_type = _da_attr_type(index, embedding_field) # get framework-specific distance and top_k function - distance_fn = _get_distance_fn(embedding_type, metric) + metric_fn = _get_metric_fn(embedding_type, metric) top_k_fn = _get_topk_fn(embedding_type) # extract embeddings from query and index @@ -121,7 +197,7 @@ def find_batched( query_embeddings = _extraxt_embeddings(query, embedding_field, embedding_type) # compute distances and return top results - dists = distance_fn(query_embeddings, index_embeddings, device=device) + dists = metric_fn(query_embeddings, index_embeddings, device=device) top_scores, top_indices = top_k_fn( dists, k=limit, device=device, descending=descending ) @@ -225,15 +301,17 @@ def _get_topk_fn(embedding_type: Type) -> Callable: ) -def _get_distance_fn(embedding_type: Type, distance_name: str) -> Callable: +def _get_metric_fn(embedding_type: Type, metric: Union[str, Callable]) -> Callable: """Dynamically import the distance function from the framework-specific module. :param embedding_type: the type of the embedding - :param distance_name: the name of the distance function - :return: the framework-specific distance function + :param metric: the name of the metric, or the metric itself + :return: the framework-specific metric """ + if callable(metric): + return metric framework = type_to_framework[embedding_type] return getattr( importlib.import_module(f'docarray.utility.math.metrics.{framework}'), - f'{distance_name}', + f'{metric}', ) From 6e48cdbf4ac9f097020ba3ae67a71c600d8c1543 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Wed, 14 Dec 2022 10:23:16 +0100 Subject: [PATCH 15/17] fix: mypy and some comments Signed-off-by: Johannes Messner --- docarray/typing/tensor/__init__.py | 6 ++++-- docarray/utility/find.py | 8 ++++++-- docarray/utility/helper/torch.py | 4 ++-- docarray/utility/math/metrics/numpy.py | 20 ++++++++++---------- docarray/utility/math/metrics/torch.py | 24 ++++++++++++------------ 5 files changed, 34 insertions(+), 28 deletions(-) diff --git a/docarray/typing/tensor/__init__.py b/docarray/typing/tensor/__init__.py index abb578ff401..1a7e1f190b4 100644 --- a/docarray/typing/tensor/__init__.py +++ b/docarray/typing/tensor/__init__.py @@ -1,3 +1,5 @@ +from typing import Dict, Set, Type + from docarray.typing.tensor.embedding import Embedding, NdArrayEmbedding from docarray.typing.tensor.ndarray import NdArray from docarray.typing.tensor.tensor import Tensor @@ -11,7 +13,7 @@ 'type_to_framework', ] -framework_types = {'numpy': {NdArray, NdArrayEmbedding}, 'torch': set()} +framework_types: Dict[str, Set] = {'numpy': {NdArray, NdArrayEmbedding}, 'torch': set()} try: import torch # noqa: F401 @@ -24,7 +26,7 @@ __all__.extend(['TorchEmbedding', 'TorchTensor']) framework_types['torch'] = {TorchTensor, TorchEmbedding} -type_to_framework = { +type_to_framework: Dict[Type, str] = { type_: framework for framework, type_set in framework_types.items() for type_ in type_set diff --git a/docarray/utility/find.py b/docarray/utility/find.py index 32d58553bee..53df66eb877 100644 --- a/docarray/utility/find.py +++ b/docarray/utility/find.py @@ -1,12 +1,14 @@ import importlib from typing import Callable, List, NamedTuple, Optional, Type, Union -import torch +import torch # TODO(johannes) this breaks the optional import of torch from docarray import Document, DocumentArray from docarray.typing import Tensor from docarray.typing.tensor import type_to_framework +# but will be fixed once we have a computational backend + class FindResult(NamedTuple): documents: DocumentArray @@ -106,7 +108,7 @@ def find_batched( index: DocumentArray, query: Union[Tensor, DocumentArray], embedding_field: str = 'embedding', - metric: Union[str, Callable[['Tensor', 'Tensor'], 'Tensor']] = 'cosine_sim', + metric: str = 'cosine_sim', limit: int = 10, device: Optional[str] = None, descending: Optional[bool] = None, @@ -289,6 +291,7 @@ def _da_attr_type(da: DocumentArray, attr: str) -> Type: def _get_topk_fn(embedding_type: Type) -> Callable: """Dynamically import the distance function from the framework-specific module. + This will go away once we have a computational backend. :param embedding_type: the type of the embedding :param distance_name: the name of the distance function @@ -303,6 +306,7 @@ def _get_topk_fn(embedding_type: Type) -> Callable: def _get_metric_fn(embedding_type: Type, metric: Union[str, Callable]) -> Callable: """Dynamically import the distance function from the framework-specific module. + This will go away once we have a proper computational backend. :param embedding_type: the type of the embedding :param metric: the name of the metric, or the metric itself diff --git a/docarray/utility/helper/torch.py b/docarray/utility/helper/torch.py index adf1fe8f28c..1e5bb92722a 100644 --- a/docarray/utility/helper/torch.py +++ b/docarray/utility/helper/torch.py @@ -4,11 +4,11 @@ def top_k( - values: 'torch.tensor', + values: torch.Tensor, k: int, descending: bool = False, device: Optional[str] = None, -) -> Tuple['torch.tensor', 'torch.tensor']: +) -> Tuple[torch.Tensor, torch.Tensor]: """ Retrieves the top k smallest values in `values`, and returns them alongside their indices in the input `values`. diff --git a/docarray/utility/math/metrics/numpy.py b/docarray/utility/math/metrics/numpy.py index b9e1810b4d4..0efe5ca40da 100644 --- a/docarray/utility/math/metrics/numpy.py +++ b/docarray/utility/math/metrics/numpy.py @@ -4,7 +4,7 @@ import numpy as np -def _expand_if_single_axis(*matrices: 'np.array') -> List['np.array']: +def _expand_if_single_axis(*matrices: np.ndarray) -> List[np.ndarray]: """Expands arrays that only have one axis, at dim 0. This ensures that all outputs can be treated as matrices, not vectors. @@ -21,18 +21,18 @@ def _expand_if_single_axis(*matrices: 'np.array') -> List['np.array']: return expanded -def _expand_if_scalar(arr: 'np.array') -> 'np.array': +def _expand_if_scalar(arr: np.ndarray) -> np.ndarray: if len(arr.shape) == 0: # avoid scalar output arr = np.expand_dims(arr, axis=0) return arr def cosine_sim( - x_mat: 'np.ndarray', - y_mat: 'np.ndarray', + x_mat: np.ndarray, + y_mat: np.ndarray, eps: float = 1e-7, device: Optional[str] = None, -) -> 'np.ndarray': +) -> np.ndarray: """Pairwise cosine similarities between all vectors in x_mat and y_mat. :param x_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is the @@ -63,11 +63,11 @@ def cosine_sim( def sqeuclidean_dist( - x_mat: 'np.ndarray', - y_mat: 'np.ndarray', + x_mat: np.ndarray, + y_mat: np.ndarray, eps: float = 1e-7, device: Optional[str] = None, -) -> 'np.ndarray': +) -> np.ndarray: """Pairwise Squared Euclidian distances between all vectors in x_mat and y_mat. :param x_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is the @@ -100,8 +100,8 @@ def sqeuclidean_dist( def euclidean_dist( - x_mat: 'np.ndarray', y_mat: 'np.ndarray', device: Optional[str] = None -) -> 'np.ndarray': + x_mat: np.ndarray, y_mat: np.ndarray, device: Optional[str] = None +) -> np.ndarray: """Pairwise Euclidian distances between all vectors in x_mat and y_mat. :param x_mat: np.ndarray of shape (n_vectors, n_dim), where n_vectors is the diff --git a/docarray/utility/math/metrics/torch.py b/docarray/utility/math/metrics/torch.py index f0b3942e832..2e1b12a50f8 100644 --- a/docarray/utility/math/metrics/torch.py +++ b/docarray/utility/math/metrics/torch.py @@ -1,12 +1,9 @@ -from typing import TYPE_CHECKING, List, Optional +from typing import List, Optional import torch -if TYPE_CHECKING: # pragma: no cover - from torch import tensor - -def _unsqueeze_if_single_axis(*matrices: 'tensor') -> List['tensor']: +def _unsqueeze_if_single_axis(*matrices: torch.Tensor) -> List[torch.Tensor]: """Unsqueezes tensors that only have one axis, at dim 0. This ensures that all outputs can be treated as matrices, not vectors. @@ -23,15 +20,18 @@ def _unsqueeze_if_single_axis(*matrices: 'tensor') -> List['tensor']: return unsqueezed -def _usqueeze_if_scalar(t: 'tensor'): +def _usqueeze_if_scalar(t: torch.Tensor): if len(t.shape) == 0: # avoid scalar output t = t.unsqueeze(0) return t def cosine_sim( - x_mat: 'tensor', y_mat: 'tensor', eps: float = 1e-7, device: Optional[str] = None -) -> 'tensor': + x_mat: torch.Tensor, + y_mat: torch.Tensor, + eps: float = 1e-7, + device: Optional[str] = None, +) -> torch.Tensor: """Pairwise cosine similarities between all vectors in x_mat and y_mat. :param x_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the @@ -61,8 +61,8 @@ def cosine_sim( def euclidean_dist( - x_mat: 'tensor', y_mat: 'tensor', device: Optional[str] = None -) -> 'tensor': + x_mat: torch.Tensor, y_mat: torch.Tensor, device: Optional[str] = None +) -> torch.Tensor: """Pairwise Euclidian distances between all vectors in x_mat and y_mat. :param x_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the @@ -88,8 +88,8 @@ def euclidean_dist( def sqeuclidean_dist( - x_mat: 'tensor', y_mat: 'tensor', device: Optional[str] = None -) -> 'tensor': + x_mat: torch.Tensor, y_mat: torch.Tensor, device: Optional[str] = None +) -> torch.Tensor: """Pairwise Squared Euclidian distances between all vectors in x_mat and y_mat. :param x_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the From b5deadc11ac4f717b0cd2c1ea5a1b4248e89f09f Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Wed, 14 Dec 2022 12:13:41 +0100 Subject: [PATCH 16/17] refactor: address code review Signed-off-by: Johannes Messner --- docarray/utility/distance.py | 1 - docarray/utility/find.py | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) delete mode 100644 docarray/utility/distance.py diff --git a/docarray/utility/distance.py b/docarray/utility/distance.py deleted file mode 100644 index 8b137891791..00000000000 --- a/docarray/utility/distance.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/docarray/utility/find.py b/docarray/utility/find.py index 53df66eb877..fe2b32d3530 100644 --- a/docarray/utility/find.py +++ b/docarray/utility/find.py @@ -204,12 +204,12 @@ class MyDocument(Document): dists, k=limit, device=device, descending=descending ) + index_doc_type = index.document_type results = [] for indices_per_query, scores_per_query in zip(top_indices, top_scores): - docs_per_query = [] + docs_per_query = DocumentArray[index_doc_type]([]) for idx in indices_per_query: # workaround until #930 is fixed docs_per_query.append(index[idx]) - docs_per_query = DocumentArray(docs_per_query) results.append(FindResult(scores=scores_per_query, documents=docs_per_query)) return results From 73e633e5e292ad311c680db48195f84edb2d91da Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Wed, 14 Dec 2022 14:18:32 +0100 Subject: [PATCH 17/17] fix: suppress mypy Signed-off-by: Johannes Messner --- docarray/utility/find.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docarray/utility/find.py b/docarray/utility/find.py index fe2b32d3530..dcf57395d4d 100644 --- a/docarray/utility/find.py +++ b/docarray/utility/find.py @@ -207,7 +207,7 @@ class MyDocument(Document): index_doc_type = index.document_type results = [] for indices_per_query, scores_per_query in zip(top_indices, top_scores): - docs_per_query = DocumentArray[index_doc_type]([]) + docs_per_query = DocumentArray[index_doc_type]([]) # type: ignore for idx in indices_per_query: # workaround until #930 is fixed docs_per_query.append(index[idx]) results.append(FindResult(scores=scores_per_query, documents=docs_per_query))