From 1addc114cd3ae55f974e0a879b81588dec658d24 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 29 Nov 2022 11:13:57 +0100 Subject: [PATCH 01/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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 f28cf79998df3f8167b7d3606daa11094b03b3ed Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Wed, 14 Dec 2022 12:05:50 +0100 Subject: [PATCH 16/23] refactor: add computational backends Signed-off-by: Johannes Messner --- docarray/computation/__init__.py | 3 + docarray/computation/abstract_comp_backend.py | 117 ++++++++++ docarray/computation/numpy_backend.py | 200 ++++++++++++++++++ docarray/computation/torch_backend.py | 167 +++++++++++++++ 4 files changed, 487 insertions(+) create mode 100644 docarray/computation/__init__.py create mode 100644 docarray/computation/abstract_comp_backend.py create mode 100644 docarray/computation/numpy_backend.py create mode 100644 docarray/computation/torch_backend.py diff --git a/docarray/computation/__init__.py b/docarray/computation/__init__.py new file mode 100644 index 00000000000..570505565c6 --- /dev/null +++ b/docarray/computation/__init__.py @@ -0,0 +1,3 @@ +from docarray.computation.abstract_comp_backend import AbstractComputationalBackend + +__all__ = ['AbstractComputationalBackend'] diff --git a/docarray/computation/abstract_comp_backend.py b/docarray/computation/abstract_comp_backend.py new file mode 100644 index 00000000000..d61bc15b61e --- /dev/null +++ b/docarray/computation/abstract_comp_backend.py @@ -0,0 +1,117 @@ +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Optional, Tuple + +if TYPE_CHECKING: + from docarray.typing.tensor.tensor import Tensor + + +class AbstractComputationalBackend(ABC): + """ + Abstract base class for computational backends. + Every supported tensor/ML framework (numpy, torch etc.) should define its own + computational backend exposing common functionality expressed in that framework. + That way, DocArray can leverage native implementations from all frameworks. + """ + + class Retrieval(ABC): + """ + Abstract class for retrieval and ranking functionalities + """ + + @staticmethod + @abstractmethod + def top_k( + values: 'Tensor', + k: int, + descending: bool = False, + device: Optional[str] = None, + ) -> Tuple['Tensor', '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 to True. + + :param values: Tensor of values to rank. + 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. + :return: Tuple containing the retrieved values, and their indices. + Both ar of shape (n_queries, k) + """ + ... + + class Metrics(ABC): + """ + Abstract base class for metrics (distances and similarities). + """ + + @staticmethod + @abstractmethod + def cosine_sim( + x_mat: 'Tensor', + y_mat: 'Tensor', + eps: float = 1e-7, + device: Optional[str] = None, + ) -> '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. + :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 computations. + If not provided, the devices of x_mat and y_mat are used. + :return: Tensor 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]. + """ + ... + + @staticmethod + @abstractmethod + 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. + + :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. + If not provided, the devices of x_mat and y_mat are used. + :return: Tensor 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]. + """ + ... + + @staticmethod + @abstractmethod + 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. + + :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. + If not provided, the devices of x_mat and y_mat are used. + :return: Tensor 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]. + """ + ... diff --git a/docarray/computation/numpy_backend.py b/docarray/computation/numpy_backend.py new file mode 100644 index 00000000000..715321df09d --- /dev/null +++ b/docarray/computation/numpy_backend.py @@ -0,0 +1,200 @@ +import warnings +from typing import List, Optional, Tuple + +import numpy as np + +from docarray.computation import AbstractComputationalBackend + + +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. + + :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 _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 + + +class NumpyCompBackend(AbstractComputationalBackend): + """ + Computational backend for Numpy. + """ + + class Retrieval: + """ + Abstract class for retrieval and ranking functionalities + """ + + @staticmethod + 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 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. + Both ar of shape (n_queries, 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 + + 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 + + class Metrics: + """ + Abstract base class for metrics (distances and similarities). + """ + + @staticmethod + def cosine_sim( + x_mat: np.ndarray, + y_mat: np.ndarray, + eps: float = 1e-7, + device: Optional[str] = None, + ) -> 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 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') + + x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat) + + 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) + + @classmethod + def euclidean_dist( + cls, 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') + + x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat) + + return _expand_if_scalar( + np.sqrt(cls.sqeuclidean_dist(x_mat, y_mat)).squeeze() + ) + + @staticmethod + 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. + + :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 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. + 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') + + x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat) + + 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) diff --git a/docarray/computation/torch_backend.py b/docarray/computation/torch_backend.py new file mode 100644 index 00000000000..e37107685d6 --- /dev/null +++ b/docarray/computation/torch_backend.py @@ -0,0 +1,167 @@ +from typing import List, Optional, Tuple + +import torch + +from docarray.computation.abstract_comp_backend import AbstractComputationalBackend + + +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. + + :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 _usqueeze_if_scalar(t: torch.Tensor): + if len(t.shape) == 0: # avoid scalar output + t = t.unsqueeze(0) + return t + + +class TorchCompBackend(AbstractComputationalBackend): + """ + Computational backend for PyTorch. + """ + + class Retrieval: + """ + Abstract class for retrieval and ranking functionalities + """ + + @staticmethod + 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 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. + 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 + ) + + class Metrics: + """ + Abstract base class for metrics (distances and similarities). + """ + + @staticmethod + def cosine_sim( + 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 + 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: torch Tensor 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) + sims = torch.mm(a_norm, b_norm.transpose(0, 1)).squeeze() + return _usqueeze_if_scalar(sims) + + @staticmethod + def euclidean_dist( + 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 + 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: torch Tensor 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) + + dists = torch.cdist(x_mat, y_mat).squeeze() + return _usqueeze_if_scalar(dists) + + @staticmethod + def sqeuclidean_dist( + 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 + 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: torch Tensor 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 _usqueeze_if_scalar((torch.cdist(x_mat, y_mat) ** 2).squeeze()) From 0177836247d6a3a8cfb21301b1a76c961ba423ad Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Wed, 14 Dec 2022 13:09:00 +0100 Subject: [PATCH 17/23] refactor: use comp backends in find function Signed-off-by: Johannes Messner --- docarray/typing/abstract_type.py | 2 + docarray/typing/tensor/__init__.py | 13 ---- docarray/typing/tensor/abstract_tensor.py | 8 +++ docarray/typing/tensor/ndarray.py | 8 +++ docarray/typing/tensor/torch_tensor.py | 8 +++ docarray/utility/find.py | 82 +++++------------------ 6 files changed, 44 insertions(+), 77 deletions(-) diff --git a/docarray/typing/abstract_type.py b/docarray/typing/abstract_type.py index 78c14d596c0..d54ec3eb585 100644 --- a/docarray/typing/abstract_type.py +++ b/docarray/typing/abstract_type.py @@ -28,3 +28,5 @@ def validate( @abstractmethod def from_protobuf(cls: Type[T], pb_msg: T) -> T: ... + + is_tensor = False diff --git a/docarray/typing/tensor/__init__.py b/docarray/typing/tensor/__init__.py index 1a7e1f190b4..58568a08c36 100644 --- a/docarray/typing/tensor/__init__.py +++ b/docarray/typing/tensor/__init__.py @@ -1,5 +1,3 @@ -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 @@ -9,12 +7,8 @@ 'Tensor', 'Embedding', 'NdArrayEmbedding', - 'framework_types', - 'type_to_framework', ] -framework_types: Dict[str, Set] = {'numpy': {NdArray, NdArrayEmbedding}, 'torch': set()} - try: import torch # noqa: F401 except ImportError: @@ -24,10 +18,3 @@ from docarray.typing.tensor.torch_tensor import TorchTensor # noqa: F401 __all__.extend(['TorchEmbedding', 'TorchTensor']) - framework_types['torch'] = {TorchTensor, TorchEmbedding} - -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/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py index 8801854ecef..8ab6e86dbc8 100644 --- a/docarray/typing/tensor/abstract_tensor.py +++ b/docarray/typing/tensor/abstract_tensor.py @@ -2,6 +2,7 @@ from abc import ABC from typing import TYPE_CHECKING, Any, Generic, List, Tuple, Type, TypeVar, Union +from docarray.computation import AbstractComputationalBackend from docarray.typing.abstract_type import AbstractType if TYPE_CHECKING: @@ -17,6 +18,7 @@ class AbstractTensor(AbstractType, Generic[ShapeT], ABC): __parametrized_meta__ = type + is_tensor = True @classmethod @abc.abstractmethod @@ -105,3 +107,9 @@ def to_protobuf(self) -> 'NdArrayProto': """ transform self into a NdArrayProto protobuf message """ + + @staticmethod + @abc.abstractmethod + def get_comp_backend() -> Type[AbstractComputationalBackend]: + """The computational backend compatible with this tensor type.""" + ... diff --git a/docarray/typing/tensor/ndarray.py b/docarray/typing/tensor/ndarray.py index 14e7174dd29..5c7d8115aff 100644 --- a/docarray/typing/tensor/ndarray.py +++ b/docarray/typing/tensor/ndarray.py @@ -20,6 +20,7 @@ from pydantic import BaseConfig from pydantic.fields import ModelField + from docarray.computation.numpy_backend import NumpyCompBackend from docarray.proto import NdArrayProto, NodeProto T = TypeVar('T', bound='NdArray') @@ -207,3 +208,10 @@ def __docarray_stack__( ) -> T: """Stack a sequence of ndarrays into a single ndarray.""" return cls.from_ndarray(np.stack(seq)) + + @staticmethod + def get_comp_backend() -> Type['NumpyCompBackend']: + """Return the computational backend of the tensor""" + from docarray.computation.numpy_backend import NumpyCompBackend + + return NumpyCompBackend diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py index 500a0de5f31..516bfef164c 100644 --- a/docarray/typing/tensor/torch_tensor.py +++ b/docarray/typing/tensor/torch_tensor.py @@ -23,6 +23,7 @@ from pydantic import BaseConfig import numpy as np from docarray.proto import NdArrayProto, NodeProto + from docarray.computation.torch_backend import TorchCompBackend from docarray.document.base_node import BaseNode @@ -239,3 +240,10 @@ def __docarray_stack__( ) -> T: """Stack a sequence of ndarray into a single ndarray.""" return cls.from_native_torch_tensor(torch.stack(seq)) # type: ignore + + @staticmethod + def get_comp_backend() -> Type['TorchCompBackend']: + """Return the computational backend of the tensor""" + from docarray.computation.torch_backend import TorchCompBackend + + return TorchCompBackend diff --git a/docarray/utility/find.py b/docarray/utility/find.py index 53df66eb877..918d1401c12 100644 --- a/docarray/utility/find.py +++ b/docarray/utility/find.py @@ -1,13 +1,7 @@ -import importlib -from typing import Callable, List, NamedTuple, Optional, Type, Union - -import torch # TODO(johannes) this breaks the optional import of torch +from typing import List, NamedTuple, Optional, Type, Union 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): @@ -189,18 +183,16 @@ class MyDocument(Document): descending = metric.endswith('_sim') # similarity metrics are descending embedding_type = _da_attr_type(index, embedding_field) - - # get framework-specific distance and top_k function - metric_fn = _get_metric_fn(embedding_type, metric) - top_k_fn = _get_topk_fn(embedding_type) + comp_backend = embedding_type.get_comp_backend() # extract embeddings from query and index index_embeddings = _extraxt_embeddings(index, embedding_field, embedding_type) query_embeddings = _extraxt_embeddings(query, embedding_field, embedding_type) # compute distances and return top results + metric_fn = getattr(comp_backend.Metrics, metric) dists = metric_fn(query_embeddings, index_embeddings, device=device) - top_scores, top_indices = top_k_fn( + top_scores, top_indices = comp_backend.Retrieval.top_k( dists, k=limit, device=device, descending=descending ) @@ -231,15 +223,9 @@ def _extract_embedding_single( 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) + # all currently supported frameworks provide `.reshape()`. Onc this is not true + # anymore, we need to add a `.reshape()` method to the computational backend + emb = emb.reshape(1, -1) return emb @@ -266,19 +252,13 @@ def _extraxt_embeddings( 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) + # all currently supported frameworks provide `.reshape()`. Onc this is not true + # anymore, we need to add a `.reshape()` method to the computational backend + emb = emb.reshape(1, -1) return emb -def _da_attr_type(da: DocumentArray, attr: str) -> Type: +def _da_attr_type(da: DocumentArray, attr: str) -> Type[Tensor]: """Get the type of the attribute according to the Document type (schema) of the DocumentArray. @@ -286,36 +266,10 @@ def _da_attr_type(da: DocumentArray, attr: str) -> Type: :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. - 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 - :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_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 - :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'{metric}', - ) + field_type = da.document_type.__fields__[attr].type_ + if not getattr(field_type, 'is_tensor', False): + raise ValueError( + f'attribute {attr} is not a tensor-like type, ' + f'but {field_type.__class__.__name__}' + ) + return field_type From 3921ac8150a40597b3c93ae1554009ca65fa15ee Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Wed, 14 Dec 2022 15:31:27 +0100 Subject: [PATCH 18/23] refactor: clean up file structure Signed-off-by: Johannes Messner --- docarray/util/__init__.py | 3 + docarray/{utility => util}/find.py | 4 +- docarray/utility/__init__.py | 3 - docarray/utility/distance.py | 1 - docarray/utility/helper/numpy.py | 51 ------ docarray/utility/helper/torch.py | 35 ---- docarray/utility/math/metrics/numpy.py | 123 ------------- docarray/utility/math/metrics/torch.py | 114 ------------ .../units/computation_backends}/__init__.py | 0 .../backend_comparisons}/__init__.py | 0 .../backend_comparisons/test_metrics.py | 59 ++++++ .../numpy_backend}/__init__.py | 0 .../numpy_backend/test_metrics.py | 67 +++++++ .../numpy_backend/test_retrieval.py | 27 +++ .../torch_backend/__init__.py | 0 .../torch_backend/test_metrics.py | 65 +++++++ .../torch_backend/test_retrieval.py | 27 +++ tests/units/util/test_find.py | 2 +- tests/units/util/test_metrics.py | 169 ------------------ tests/units/util/test_topk.py | 49 ----- 20 files changed, 251 insertions(+), 548 deletions(-) create mode 100644 docarray/util/__init__.py rename docarray/{utility => util}/find.py (99%) delete mode 100644 docarray/utility/__init__.py delete mode 100644 docarray/utility/distance.py delete mode 100644 docarray/utility/helper/numpy.py delete mode 100644 docarray/utility/helper/torch.py delete mode 100644 docarray/utility/math/metrics/numpy.py delete mode 100644 docarray/utility/math/metrics/torch.py rename {docarray/utility/helper => tests/units/computation_backends}/__init__.py (100%) rename {docarray/utility/math => tests/units/computation_backends/backend_comparisons}/__init__.py (100%) create mode 100644 tests/units/computation_backends/backend_comparisons/test_metrics.py rename {docarray/utility/math/metrics => tests/units/computation_backends/numpy_backend}/__init__.py (100%) create mode 100644 tests/units/computation_backends/numpy_backend/test_metrics.py create mode 100644 tests/units/computation_backends/numpy_backend/test_retrieval.py create mode 100644 tests/units/computation_backends/torch_backend/__init__.py create mode 100644 tests/units/computation_backends/torch_backend/test_metrics.py create mode 100644 tests/units/computation_backends/torch_backend/test_retrieval.py delete mode 100644 tests/units/util/test_metrics.py delete mode 100644 tests/units/util/test_topk.py diff --git a/docarray/util/__init__.py b/docarray/util/__init__.py new file mode 100644 index 00000000000..1e44ed20dba --- /dev/null +++ b/docarray/util/__init__.py @@ -0,0 +1,3 @@ +from docarray.util.find import find, find_batched + +__all__ = ['find', 'find_batched'] diff --git a/docarray/utility/find.py b/docarray/util/find.py similarity index 99% rename from docarray/utility/find.py rename to docarray/util/find.py index 918d1401c12..a0cd76fd0f8 100644 --- a/docarray/utility/find.py +++ b/docarray/util/find.py @@ -40,7 +40,7 @@ def find( from docarray import DocumentArray, Document from docarray.typing import TorchTensor - from docarray.utility.find import find + from docarray.util.find import find class MyDocument(Document): @@ -129,7 +129,7 @@ def find_batched( from docarray import DocumentArray, Document from docarray.typing import TorchTensor - from docarray.utility.find import find + from docarray.util.find import find class MyDocument(Document): diff --git a/docarray/utility/__init__.py b/docarray/utility/__init__.py deleted file mode 100644 index 7c9bada7262..00000000000 --- a/docarray/utility/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from docarray.utility.find import find, find_batched - -__all__ = ['find', 'find_batched'] 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/helper/numpy.py b/docarray/utility/helper/numpy.py deleted file mode 100644 index cf2dcadc222..00000000000 --- a/docarray/utility/helper/numpy.py +++ /dev/null @@ -1,51 +0,0 @@ -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 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. - Both ar of shape (n_queries, 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 - - 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 deleted file mode 100644 index 1e5bb92722a..00000000000 --- a/docarray/utility/helper/torch.py +++ /dev/null @@ -1,35 +0,0 @@ -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 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. - 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/docarray/utility/math/metrics/numpy.py b/docarray/utility/math/metrics/numpy.py deleted file mode 100644 index 0efe5ca40da..00000000000 --- a/docarray/utility/math/metrics/numpy.py +++ /dev/null @@ -1,123 +0,0 @@ -import warnings -from typing import List, Optional - -import numpy as np - - -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. - - :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 _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, - eps: float = 1e-7, - device: Optional[str] = None, -) -> 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 - 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') - - x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat) - - 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_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. - - :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 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. - 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') - - x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat) - - 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_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. - - :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') - - x_mat, y_mat = _expand_if_single_axis(x_mat, y_mat) - - return _expand_if_scalar(np.sqrt(sqeuclidean_dist(x_mat, y_mat)).squeeze()) diff --git a/docarray/utility/math/metrics/torch.py b/docarray/utility/math/metrics/torch.py deleted file mode 100644 index 2e1b12a50f8..00000000000 --- a/docarray/utility/math/metrics/torch.py +++ /dev/null @@ -1,114 +0,0 @@ -from typing import List, Optional - -import torch - - -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. - - :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 _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: 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 - 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) - sims = torch.mm(a_norm, b_norm.transpose(0, 1)).squeeze() - return _usqueeze_if_scalar(sims) - - -def euclidean_dist( - 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 - 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) - - dists = torch.cdist(x_mat, y_mat).squeeze() - return _usqueeze_if_scalar(dists) - - -def sqeuclidean_dist( - 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 - 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 _usqueeze_if_scalar((torch.cdist(x_mat, y_mat) ** 2).squeeze()) diff --git a/docarray/utility/helper/__init__.py b/tests/units/computation_backends/__init__.py similarity index 100% rename from docarray/utility/helper/__init__.py rename to tests/units/computation_backends/__init__.py diff --git a/docarray/utility/math/__init__.py b/tests/units/computation_backends/backend_comparisons/__init__.py similarity index 100% rename from docarray/utility/math/__init__.py rename to tests/units/computation_backends/backend_comparisons/__init__.py diff --git a/tests/units/computation_backends/backend_comparisons/test_metrics.py b/tests/units/computation_backends/backend_comparisons/test_metrics.py new file mode 100644 index 00000000000..2db6df3eba8 --- /dev/null +++ b/tests/units/computation_backends/backend_comparisons/test_metrics.py @@ -0,0 +1,59 @@ +import torch + +from docarray.computation.numpy_backend import NumpyCompBackend +from docarray.computation.torch_backend import TorchCompBackend + +np_metrics = NumpyCompBackend.Metrics +torch_metrics = TorchCompBackend.Metrics + + +def test_cosine_sim_compare(): + a = torch.rand(128) + b = torch.rand(128) + torch.testing.assert_close( + torch_metrics.cosine_sim(a, b), + torch.from_numpy(np_metrics.cosine_sim(a.numpy(), b.numpy())), + ) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + torch.testing.assert_close( + torch_metrics.cosine_sim(a, b), + torch.from_numpy(np_metrics.cosine_sim(a.numpy(), b.numpy())), + ) + + +def test_euclidean_dist_compare(): + a = torch.rand(128) + b = torch.rand(128) + torch.testing.assert_close( + torch_metrics.euclidean_dist(a, b), + torch.from_numpy(np_metrics.euclidean_dist(a.numpy(), b.numpy())).to( + torch.float32 + ), + ) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + torch.testing.assert_close( + torch_metrics.euclidean_dist(a, b), + torch.from_numpy(np_metrics.euclidean_dist(a.numpy(), b.numpy())), + ) + + +def test_sqeuclidean_dist_compare(): + a = torch.rand(128) + b = torch.rand(128) + torch.testing.assert_close( + torch_metrics.sqeuclidean_dist(a, b), + torch.from_numpy(np_metrics.sqeuclidean_dist(a.numpy(), b.numpy())).to( + torch.float32 + ), + ) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + torch.testing.assert_close( + torch_metrics.sqeuclidean_dist(a, b), + torch.from_numpy(np_metrics.sqeuclidean_dist(a.numpy(), b.numpy())), + ) diff --git a/docarray/utility/math/metrics/__init__.py b/tests/units/computation_backends/numpy_backend/__init__.py similarity index 100% rename from docarray/utility/math/metrics/__init__.py rename to tests/units/computation_backends/numpy_backend/__init__.py diff --git a/tests/units/computation_backends/numpy_backend/test_metrics.py b/tests/units/computation_backends/numpy_backend/test_metrics.py new file mode 100644 index 00000000000..2cd0369e876 --- /dev/null +++ b/tests/units/computation_backends/numpy_backend/test_metrics.py @@ -0,0 +1,67 @@ +import numpy as np + +from docarray.computation.numpy_backend import NumpyCompBackend + +metrics = NumpyCompBackend.Metrics + + +def test_cosine_sim_np(): + a = np.random.rand(128) + b = np.random.rand(128) + assert metrics.cosine_sim(a, b).shape == (1,) + assert metrics.cosine_sim(a, b) == metrics.cosine_sim(b, a) + np.testing.assert_array_almost_equal(metrics.cosine_sim(a, a), np.ones((1,))) + + a = np.random.rand(10, 3) + b = np.random.rand(5, 3) + assert metrics.cosine_sim(a, b).shape == (10, 5) + assert metrics.cosine_sim(b, a).shape == (5, 10) + diag_dists = np.diagonal(metrics.cosine_sim(b, b)) # self-comparisons + np.testing.assert_array_almost_equal(diag_dists, np.ones((5,))) + + +def test_euclidean_dist_np(): + a = np.random.rand(128) + b = np.random.rand(128) + assert metrics.euclidean_dist(a, b).shape == (1,) + assert metrics.euclidean_dist(a, b) == metrics.euclidean_dist(b, a) + np.testing.assert_array_almost_equal(metrics.euclidean_dist(a, a), np.zeros((1,))) + + a = np.random.rand(10, 3) + b = np.random.rand(5, 3) + assert metrics.euclidean_dist(a, b).shape == (10, 5) + assert metrics.euclidean_dist(b, a).shape == (5, 10) + diag_dists = np.diagonal(metrics.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( + metrics.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( + metrics.euclidean_dist(a, b), desired_output_singleton + ) + + +def test_sqeuclidea_dist_np(): + a = np.random.rand(128) + b = np.random.rand(128) + assert metrics.sqeuclidean_dist(a, b).shape == (1,) + np.testing.assert_array_almost_equal( + metrics.sqeuclidean_dist(a, b), metrics.euclidean_dist(a, b) ** 2 + ) + + a = np.random.rand(10, 3) + b = np.random.rand(5, 3) + assert metrics.sqeuclidean_dist(a, b).shape == (10, 5) + np.testing.assert_array_almost_equal( + metrics.sqeuclidean_dist(a, b), metrics.euclidean_dist(a, b) ** 2 + ) diff --git a/tests/units/computation_backends/numpy_backend/test_retrieval.py b/tests/units/computation_backends/numpy_backend/test_retrieval.py new file mode 100644 index 00000000000..a00f254508c --- /dev/null +++ b/tests/units/computation_backends/numpy_backend/test_retrieval.py @@ -0,0 +1,27 @@ +import numpy as np + +from docarray.computation.numpy_backend import NumpyCompBackend + + +def test_topk_numpy(): + top_k = NumpyCompBackend.Retrieval.top_k + + a = np.array([1, 4, 2, 7, 4, 9, 2]) + vals, indices = 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 = 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() diff --git a/tests/units/computation_backends/torch_backend/__init__.py b/tests/units/computation_backends/torch_backend/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/units/computation_backends/torch_backend/test_metrics.py b/tests/units/computation_backends/torch_backend/test_metrics.py new file mode 100644 index 00000000000..c353e4f77bd --- /dev/null +++ b/tests/units/computation_backends/torch_backend/test_metrics.py @@ -0,0 +1,65 @@ +import torch + +from docarray.computation.torch_backend import TorchCompBackend + +metrics = TorchCompBackend.Metrics + + +def test_cosine_sim_torch(): + a = torch.rand(128) + b = torch.rand(128) + assert metrics.cosine_sim(a, b).shape == (1,) + assert metrics.cosine_sim(a, b) == metrics.cosine_sim(b, a) + torch.testing.assert_close(metrics.cosine_sim(a, a), torch.ones(1)) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + assert metrics.cosine_sim(a, b).shape == (10, 5) + assert metrics.cosine_sim(b, a).shape == (5, 10) + diag_dists = torch.diagonal(metrics.cosine_sim(b, b)) # self-comparisons + torch.testing.assert_allclose(diag_dists, torch.ones(5)) + + +def test_euclidean_dist_torch(): + a = torch.rand(128) + b = torch.rand(128) + assert metrics.euclidean_dist(a, b).shape == (1,) + assert metrics.euclidean_dist(a, b) == metrics.euclidean_dist(b, a) + torch.testing.assert_close(metrics.euclidean_dist(a, a), torch.zeros(1)) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + assert metrics.euclidean_dist(a, b).shape == (10, 5) + assert metrics.euclidean_dist(b, a).shape == (5, 10) + diag_dists = torch.diagonal(metrics.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(metrics.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(metrics.euclidean_dist(a, b), desired_output_singleton) + + +def test_sqeuclidean_dist_torch(): + a = torch.rand(128) + b = torch.rand(128) + assert metrics.sqeuclidean_dist(a, b).shape == (1,) + torch.testing.assert_close( + metrics.sqeuclidean_dist(a, b), + metrics.euclidean_dist(a, b) ** 2, + ) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + assert metrics.sqeuclidean_dist(a, b).shape == (10, 5) + torch.testing.assert_close( + metrics.sqeuclidean_dist(a, b), + metrics.euclidean_dist(a, b) ** 2, + ) diff --git a/tests/units/computation_backends/torch_backend/test_retrieval.py b/tests/units/computation_backends/torch_backend/test_retrieval.py new file mode 100644 index 00000000000..2e3a1833ff0 --- /dev/null +++ b/tests/units/computation_backends/torch_backend/test_retrieval.py @@ -0,0 +1,27 @@ +import torch + +from docarray.computation.torch_backend import TorchCompBackend + + +def test_topk(): + top_k = TorchCompBackend.Retrieval.top_k + + a = torch.tensor([1, 4, 2, 7, 4, 9, 2]) + vals, indices = 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 = 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() diff --git a/tests/units/util/test_find.py b/tests/units/util/test_find.py index 665b880912c..8ec3ccfef6c 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, find_batched +from docarray.util import find, find_batched class TorchDoc(Document): diff --git a/tests/units/util/test_metrics.py b/tests/units/util/test_metrics.py deleted file mode 100644 index c50af26fc76..00000000000 --- a/tests/units/util/test_metrics.py +++ /dev/null @@ -1,169 +0,0 @@ -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 - ) diff --git a/tests/units/util/test_topk.py b/tests/units/util/test_topk.py deleted file mode 100644 index 2bd8d3b520a..00000000000 --- a/tests/units/util/test_topk.py +++ /dev/null @@ -1,49 +0,0 @@ -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 4597115e67b6406cea3395c4420f60f6d20a8d08 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Wed, 14 Dec 2022 16:39:36 +0100 Subject: [PATCH 19/23] fix: some typing issues Signed-off-by: Johannes Messner --- docarray/computation/abstract_comp_backend.py | 12 +++++++- docarray/computation/numpy_backend.py | 8 ++++- docarray/computation/torch_backend.py | 8 ++++- docarray/typing/tensor/abstract_tensor.py | 11 ++++++- docarray/typing/tensor/ndarray.py | 21 +++++-------- docarray/typing/tensor/torch_tensor.py | 30 ++++--------------- .../document/proto/test_proto_based_object.py | 4 +-- 7 files changed, 50 insertions(+), 44 deletions(-) diff --git a/docarray/computation/abstract_comp_backend.py b/docarray/computation/abstract_comp_backend.py index d61bc15b61e..c713927f8e6 100644 --- a/docarray/computation/abstract_comp_backend.py +++ b/docarray/computation/abstract_comp_backend.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Optional, Tuple +from typing import TYPE_CHECKING, List, Optional, Tuple, Union if TYPE_CHECKING: from docarray.typing.tensor.tensor import Tensor @@ -13,6 +13,16 @@ class AbstractComputationalBackend(ABC): That way, DocArray can leverage native implementations from all frameworks. """ + @staticmethod + @abstractmethod + def stack( + tensors: Union[List['Tensor'], Tuple['Tensor']], dim: int = 0 + ) -> 'Tensor': + """ + Stack a list of tensors along a new axis. + """ + ... + class Retrieval(ABC): """ Abstract class for retrieval and ranking functionalities diff --git a/docarray/computation/numpy_backend.py b/docarray/computation/numpy_backend.py index 715321df09d..edf341a034e 100644 --- a/docarray/computation/numpy_backend.py +++ b/docarray/computation/numpy_backend.py @@ -1,5 +1,5 @@ import warnings -from typing import List, Optional, Tuple +from typing import List, Optional, Tuple, Union import numpy as np @@ -34,6 +34,12 @@ class NumpyCompBackend(AbstractComputationalBackend): Computational backend for Numpy. """ + @staticmethod + def stack( + tensors: Union[List['np.ndarray'], Tuple['np.ndarray']], dim: int = 0 + ) -> 'np.ndarray': + return np.stack(tensors, axis=dim) + class Retrieval: """ Abstract class for retrieval and ranking functionalities diff --git a/docarray/computation/torch_backend.py b/docarray/computation/torch_backend.py index e37107685d6..9bf6a945717 100644 --- a/docarray/computation/torch_backend.py +++ b/docarray/computation/torch_backend.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Tuple +from typing import List, Optional, Tuple, Union import torch @@ -33,6 +33,12 @@ class TorchCompBackend(AbstractComputationalBackend): Computational backend for PyTorch. """ + @staticmethod + def stack( + tensors: Union[List['torch.Tensor'], Tuple['torch.Tensor']], dim: int = 0 + ) -> 'torch.Tensor': + return torch.stack(tensors, dim=dim) + class Retrieval: """ Abstract class for retrieval and ranking functionalities diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py index 8ab6e86dbc8..ebfed57fffb 100644 --- a/docarray/typing/tensor/abstract_tensor.py +++ b/docarray/typing/tensor/abstract_tensor.py @@ -98,9 +98,18 @@ def __class_getitem__(cls, item: Any): return cls._docarray_create_parametrized_type(target_shape) @classmethod - @abc.abstractmethod def __docarray_stack__(cls: Type[T], seq: Union[List[T], Tuple[T]]) -> T: """Stack a sequence of tensors into a single tensor.""" + comp_backend = cls.get_comp_backend() + return cls.__docarray_from_native__(comp_backend.stack(seq)) + + @classmethod + @abc.abstractmethod + def __docarray_from_native__(cls: Type[T], value: Any) -> T: + """ + Create a DocArray tensor from a tensor that is native to the given framework, + e.g. from numpy.ndarray or torch.Tensor. + """ ... def to_protobuf(self) -> 'NdArrayProto': diff --git a/docarray/typing/tensor/ndarray.py b/docarray/typing/tensor/ndarray.py index 5c7d8115aff..0e2b489d251 100644 --- a/docarray/typing/tensor/ndarray.py +++ b/docarray/typing/tensor/ndarray.py @@ -87,7 +87,7 @@ def __docarray_validate_shape__(cls, t: T, shape: Tuple[int]) -> T: # type: ign f'of shape {t.shape} to shape {shape}' ) try: - value = cls.from_ndarray(np.reshape(t, shape)) + value = cls.__docarray_from_native__(np.reshape(t, shape)) return cast(T, value) except RuntimeError: raise ValueError( @@ -102,25 +102,25 @@ def validate( config: 'BaseConfig', ) -> T: if isinstance(value, np.ndarray): - return cls.from_ndarray(value) + return cls.__docarray_from_native__(value) elif isinstance(value, NdArray): return cast(T, value) elif isinstance(value, list) or isinstance(value, tuple): try: arr_from_list: np.ndarray = np.asarray(value) - return cls.from_ndarray(arr_from_list) + return cls.__docarray_from_native__(arr_from_list) except Exception: pass # handled below else: try: arr: np.ndarray = np.ndarray(value) - return cls.from_ndarray(arr) + return cls.__docarray_from_native__(arr) except Exception: pass # handled below raise ValueError(f'Expected a numpy.ndarray compatible type, got {type(value)}') @classmethod - def from_ndarray(cls: Type[T], value: np.ndarray) -> T: + def __docarray_from_native__(cls: Type[T], value: np.ndarray) -> T: return value.view(cls) @classmethod @@ -181,9 +181,9 @@ def from_protobuf(cls: Type[T], pb_msg: 'NdArrayProto') -> 'T': source = pb_msg.dense if source.buffer: x = np.frombuffer(bytearray(source.buffer), dtype=source.dtype) - return cls.from_ndarray(x.reshape(source.shape)) + return cls.__docarray_from_native__(x.reshape(source.shape)) elif len(source.shape) > 0: - return cls.from_ndarray(np.zeros(source.shape)) + return cls.__docarray_from_native__(np.zeros(source.shape)) else: raise ValueError(f'proto message {pb_msg} cannot be cast to a NdArray') @@ -202,13 +202,6 @@ def to_protobuf(self) -> 'NdArrayProto': return nd_proto - @classmethod - def __docarray_stack__( - cls: Type[T], seq: Union[List['NdArray'], Tuple['NdArray']] - ) -> T: - """Stack a sequence of ndarrays into a single ndarray.""" - return cls.from_ndarray(np.stack(seq)) - @staticmethod def get_comp_backend() -> Type['NumpyCompBackend']: """Return the computational backend of the tensor""" diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py index 516bfef164c..a1a6d7b9596 100644 --- a/docarray/typing/tensor/torch_tensor.py +++ b/docarray/typing/tensor/torch_tensor.py @@ -1,17 +1,6 @@ import warnings from copy import copy -from typing import ( - TYPE_CHECKING, - Any, - Dict, - Generic, - List, - Tuple, - Type, - TypeVar, - Union, - cast, -) +from typing import TYPE_CHECKING, Any, Dict, Generic, Tuple, Type, TypeVar, Union, cast import numpy as np import torch # type: ignore @@ -104,7 +93,7 @@ def __docarray_validate_shape__(cls, t: T, shape: Tuple[int]) -> T: # type: ign f'of shape {t.shape} to shape {shape}' ) try: - value = cls.from_native_torch_tensor(t.view(shape)) + value = cls.__docarray_from_native__(t.view(shape)) return cast(T, value) except RuntimeError: raise ValueError( @@ -121,12 +110,12 @@ def validate( if isinstance(value, TorchTensor): return cast(T, value) elif isinstance(value, torch.Tensor): - return cls.from_native_torch_tensor(value) + return cls.__docarray_from_native__(value) else: try: arr: torch.Tensor = torch.tensor(value) - return cls.from_native_torch_tensor(arr) + return cls.__docarray_from_native__(arr) except Exception: pass # handled below raise ValueError(f'Expected a torch.Tensor compatible type, got {type(value)}') @@ -172,7 +161,7 @@ def unwrap(self) -> torch.Tensor: return value @classmethod - def from_native_torch_tensor(cls: Type[T], value: torch.Tensor) -> T: + def __docarray_from_native__(cls: Type[T], value: torch.Tensor) -> T: """Create a TorchTensor from a native torch.Tensor :param value: the native torch.Tensor @@ -188,7 +177,7 @@ def from_ndarray(cls: Type[T], value: np.ndarray) -> T: :param value: the numpy array :return: a TorchTensor """ - return cls.from_native_torch_tensor(torch.from_numpy(value)) + return cls.__docarray_from_native__(torch.from_numpy(value)) def _to_node_protobuf(self: T, field: str = 'torch_tensor') -> 'NodeProto': """Convert Document into a NodeProto protobuf message. This function should @@ -234,13 +223,6 @@ def to_protobuf(self) -> 'NdArrayProto': return nd_proto - @classmethod - def __docarray_stack__( - cls: Type[T], seq: Union[Tuple['TorchTensor'], List['TorchTensor']] - ) -> T: - """Stack a sequence of ndarray into a single ndarray.""" - return cls.from_native_torch_tensor(torch.stack(seq)) # type: ignore - @staticmethod def get_comp_backend() -> Type['TorchCompBackend']: """Return the computational backend of the tensor""" diff --git a/tests/units/document/proto/test_proto_based_object.py b/tests/units/document/proto/test_proto_based_object.py index e5e0c10bfa0..7930d32c395 100644 --- a/tests/units/document/proto/test_proto_based_object.py +++ b/tests/units/document/proto/test_proto_based_object.py @@ -17,7 +17,7 @@ def test_ndarray(): original_ndarray = np.zeros((3, 224, 224)) - custom_ndarray = NdArray.from_ndarray(original_ndarray) + custom_ndarray = NdArray.__docarray_from_native__(original_ndarray) tensor = NdArray.from_protobuf(custom_ndarray.to_protobuf()) @@ -30,7 +30,7 @@ def test_document_proto_set(): nested_item1 = NodeProto(text='hello') - ndarray = NdArray.from_ndarray(np.zeros((3, 224, 224))) + ndarray = NdArray.__docarray_from_native__(np.zeros((3, 224, 224))) nd_proto = ndarray.to_protobuf() nested_item2 = NodeProto(ndarray=nd_proto) From 4a6e1b3a7df84a39e59606930a7c5cb3f0501197 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Thu, 15 Dec 2022 15:20:13 +0100 Subject: [PATCH 20/23] fix: reduce mypy errors Signed-off-by: Johannes Messner --- docarray/computation/abstract_comp_backend.py | 41 +++++++++++-------- docarray/computation/numpy_backend.py | 12 +++--- docarray/computation/torch_backend.py | 10 ++--- 3 files changed, 34 insertions(+), 29 deletions(-) diff --git a/docarray/computation/abstract_comp_backend.py b/docarray/computation/abstract_comp_backend.py index c713927f8e6..f6ebc5751df 100644 --- a/docarray/computation/abstract_comp_backend.py +++ b/docarray/computation/abstract_comp_backend.py @@ -1,11 +1,14 @@ +import typing from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, List, Optional, Tuple, Union +from typing import List, Optional, Tuple, TypeVar, Union -if TYPE_CHECKING: - from docarray.typing.tensor.tensor import Tensor +# In practice all of the below will be the same type +TTensor = TypeVar('TTensor') +TTensorRetrieval = TypeVar('TTensorRetrieval') +TTensorMetrics = TypeVar('TTensorMetrics') -class AbstractComputationalBackend(ABC): +class AbstractComputationalBackend(ABC, typing.Generic[TTensor]): """ Abstract base class for computational backends. Every supported tensor/ML framework (numpy, torch etc.) should define its own @@ -16,14 +19,14 @@ class AbstractComputationalBackend(ABC): @staticmethod @abstractmethod def stack( - tensors: Union[List['Tensor'], Tuple['Tensor']], dim: int = 0 - ) -> 'Tensor': + tensors: Union[List['TTensor'], Tuple['TTensor']], dim: int = 0 + ) -> 'TTensor': """ Stack a list of tensors along a new axis. """ ... - class Retrieval(ABC): + class Retrieval(ABC, typing.Generic[TTensorRetrieval]): """ Abstract class for retrieval and ranking functionalities """ @@ -31,11 +34,11 @@ class Retrieval(ABC): @staticmethod @abstractmethod def top_k( - values: 'Tensor', + values: 'TTensorRetrieval', k: int, descending: bool = False, device: Optional[str] = None, - ) -> Tuple['Tensor', 'Tensor']: + ) -> Tuple['TTensorRetrieval', 'TTensorRetrieval']: """ Retrieves the top k smallest values in `values`, and returns them alongside their indices in the input `values`. @@ -54,7 +57,7 @@ def top_k( """ ... - class Metrics(ABC): + class Metrics(ABC, typing.Generic[TTensorMetrics]): """ Abstract base class for metrics (distances and similarities). """ @@ -62,11 +65,11 @@ class Metrics(ABC): @staticmethod @abstractmethod def cosine_sim( - x_mat: 'Tensor', - y_mat: 'Tensor', + x_mat: 'TTensorMetrics', + y_mat: 'TTensorMetrics', eps: float = 1e-7, device: Optional[str] = None, - ) -> 'Tensor': + ) -> 'TTensorMetrics': """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 @@ -86,8 +89,10 @@ def cosine_sim( @staticmethod @abstractmethod def euclidean_dist( - x_mat: 'Tensor', y_mat: 'Tensor', device: Optional[str] = None - ) -> 'Tensor': + x_mat: 'TTensorMetrics', + y_mat: 'TTensorMetrics', + device: Optional[str] = None, + ) -> 'TTensorMetrics': """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 @@ -106,8 +111,10 @@ def euclidean_dist( @staticmethod @abstractmethod def sqeuclidean_dist( - x_mat: 'Tensor', y_mat: 'Tensor', device: Optional[str] = None - ) -> 'Tensor': + x_mat: 'TTensorMetrics', + y_mat: 'TTensorMetrics', + device: Optional[str] = None, + ) -> 'TTensorMetrics': """Pairwise Squared Euclidian distances between all vectors in x_mat and y_mat. diff --git a/docarray/computation/numpy_backend.py b/docarray/computation/numpy_backend.py index edf341a034e..51e77f3cb0c 100644 --- a/docarray/computation/numpy_backend.py +++ b/docarray/computation/numpy_backend.py @@ -29,7 +29,7 @@ def _expand_if_scalar(arr: np.ndarray) -> np.ndarray: return arr -class NumpyCompBackend(AbstractComputationalBackend): +class NumpyCompBackend(AbstractComputationalBackend[np.ndarray]): """ Computational backend for Numpy. """ @@ -40,7 +40,7 @@ def stack( ) -> 'np.ndarray': return np.stack(tensors, axis=dim) - class Retrieval: + class Retrieval(AbstractComputationalBackend.Retrieval[np.ndarray]): """ Abstract class for retrieval and ranking functionalities """ @@ -92,7 +92,7 @@ def top_k( return values, idx - class Metrics: + class Metrics(AbstractComputationalBackend.Metrics[np.ndarray]): """ Abstract base class for metrics (distances and similarities). """ @@ -169,7 +169,6 @@ def euclidean_dist( 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 @@ -181,15 +180,14 @@ def sqeuclidean_dist( :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. The index [i_x, i_y] contains the cosine Squared Euclidian between x_mat[i_x] and y_mat[i_y]. """ + eps: float = 1e-7 # avoid problems with numerical inaccuracies + if device is not None: warnings.warn('`device` is not supported for numpy operations') diff --git a/docarray/computation/torch_backend.py b/docarray/computation/torch_backend.py index 9bf6a945717..e91492bd6a6 100644 --- a/docarray/computation/torch_backend.py +++ b/docarray/computation/torch_backend.py @@ -28,7 +28,7 @@ def _usqueeze_if_scalar(t: torch.Tensor): return t -class TorchCompBackend(AbstractComputationalBackend): +class TorchCompBackend(AbstractComputationalBackend[torch.Tensor]): """ Computational backend for PyTorch. """ @@ -39,18 +39,18 @@ def stack( ) -> 'torch.Tensor': return torch.stack(tensors, dim=dim) - class Retrieval: + class Retrieval(AbstractComputationalBackend.Retrieval[torch.Tensor]): """ Abstract class for retrieval and ranking functionalities """ @staticmethod 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`. @@ -78,7 +78,7 @@ def top_k( input=values, k=k, largest=descending, sorted=True, dim=-1 ) - class Metrics: + class Metrics(AbstractComputationalBackend.Metrics[torch.Tensor]): """ Abstract base class for metrics (distances and similarities). """ From 650376f3f6d367ede3d518143b649d166743263f Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Thu, 15 Dec 2022 16:57:59 +0100 Subject: [PATCH 21/23] fix: mypy Signed-off-by: Johannes Messner --- docarray/typing/tensor/abstract_tensor.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py index a757a136419..89702ae5cf6 100644 --- a/docarray/typing/tensor/abstract_tensor.py +++ b/docarray/typing/tensor/abstract_tensor.py @@ -6,8 +6,11 @@ from docarray.typing.abstract_type import AbstractType if TYPE_CHECKING: + import numpy as np + import torch from pydantic import BaseConfig from pydantic.fields import ModelField +TComp = TypeVar('TComp', bound=Union['torch.Tensor', 'np.ndarray']) T = TypeVar('T', bound='AbstractTensor') ShapeT = TypeVar('ShapeT') @@ -99,7 +102,9 @@ def __class_getitem__(cls, item: Any): def __docarray_stack__(cls: Type[T], seq: Union[List[T], Tuple[T]]) -> T: """Stack a sequence of tensors into a single tensor.""" comp_backend = cls.get_comp_backend() - return cls.__docarray_from_native__(comp_backend.stack(seq)) + # at runtime, 'T' is always the correct input type for .stack() + # but mypy doesn't know that, so we ignore it here + return cls.__docarray_from_native__(comp_backend.stack(seq)) # type: ignore @classmethod @abc.abstractmethod From 73a43580129e2e915cfd5beda30decbfc3b67b4a Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Thu, 15 Dec 2022 17:02:55 +0100 Subject: [PATCH 22/23] fix: mypy Signed-off-by: Johannes Messner --- docarray/typing/tensor/abstract_tensor.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py index 89702ae5cf6..b472ecd58ce 100644 --- a/docarray/typing/tensor/abstract_tensor.py +++ b/docarray/typing/tensor/abstract_tensor.py @@ -6,11 +6,8 @@ from docarray.typing.abstract_type import AbstractType if TYPE_CHECKING: - import numpy as np - import torch from pydantic import BaseConfig from pydantic.fields import ModelField -TComp = TypeVar('TComp', bound=Union['torch.Tensor', 'np.ndarray']) T = TypeVar('T', bound='AbstractTensor') ShapeT = TypeVar('ShapeT') From 4b50d0f2770b2bbe91bf92c4d73de9bcb989d969 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Fri, 16 Dec 2022 11:43:42 +0100 Subject: [PATCH 23/23] refactor: remove is_tensor flag Signed-off-by: Johannes Messner --- docarray/typing/abstract_type.py | 2 -- docarray/typing/tensor/abstract_tensor.py | 1 - docarray/util/find.py | 3 ++- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/docarray/typing/abstract_type.py b/docarray/typing/abstract_type.py index e248f5ce696..3d2e09c8602 100644 --- a/docarray/typing/abstract_type.py +++ b/docarray/typing/abstract_type.py @@ -13,8 +13,6 @@ class AbstractType(BaseNode): - is_tensor = False # change for tensor-like subclasses - @classmethod def __get_validators__(cls): yield cls.validate diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py index b472ecd58ce..39ec52ff3e5 100644 --- a/docarray/typing/tensor/abstract_tensor.py +++ b/docarray/typing/tensor/abstract_tensor.py @@ -16,7 +16,6 @@ class AbstractTensor(AbstractType, Generic[ShapeT], ABC): __parametrized_meta__ = type - is_tensor = True @classmethod @abc.abstractmethod diff --git a/docarray/util/find.py b/docarray/util/find.py index a0cd76fd0f8..004ce942c0e 100644 --- a/docarray/util/find.py +++ b/docarray/util/find.py @@ -2,6 +2,7 @@ from docarray import Document, DocumentArray from docarray.typing import Tensor +from docarray.typing.tensor.abstract_tensor import AbstractTensor class FindResult(NamedTuple): @@ -267,7 +268,7 @@ def _da_attr_type(da: DocumentArray, attr: str) -> Type[Tensor]: :return: the type of the attribute """ field_type = da.document_type.__fields__[attr].type_ - if not getattr(field_type, 'is_tensor', False): + if not issubclass(field_type, AbstractTensor): raise ValueError( f'attribute {attr} is not a tensor-like type, ' f'but {field_type.__class__.__name__}'