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..f6ebc5751df --- /dev/null +++ b/docarray/computation/abstract_comp_backend.py @@ -0,0 +1,134 @@ +import typing +from abc import ABC, abstractmethod +from typing import List, Optional, Tuple, TypeVar, Union + +# In practice all of the below will be the same type +TTensor = TypeVar('TTensor') +TTensorRetrieval = TypeVar('TTensorRetrieval') +TTensorMetrics = TypeVar('TTensorMetrics') + + +class AbstractComputationalBackend(ABC, typing.Generic[TTensor]): + """ + 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. + """ + + @staticmethod + @abstractmethod + def stack( + tensors: Union[List['TTensor'], Tuple['TTensor']], dim: int = 0 + ) -> 'TTensor': + """ + Stack a list of tensors along a new axis. + """ + ... + + class Retrieval(ABC, typing.Generic[TTensorRetrieval]): + """ + Abstract class for retrieval and ranking functionalities + """ + + @staticmethod + @abstractmethod + def top_k( + values: 'TTensorRetrieval', + k: int, + descending: bool = False, + device: Optional[str] = None, + ) -> Tuple['TTensorRetrieval', 'TTensorRetrieval']: + """ + 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, typing.Generic[TTensorMetrics]): + """ + Abstract base class for metrics (distances and similarities). + """ + + @staticmethod + @abstractmethod + def cosine_sim( + x_mat: 'TTensorMetrics', + y_mat: 'TTensorMetrics', + eps: float = 1e-7, + device: Optional[str] = None, + ) -> '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 + 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: '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 + 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: 'TTensorMetrics', + y_mat: 'TTensorMetrics', + device: Optional[str] = None, + ) -> 'TTensorMetrics': + """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..51e77f3cb0c --- /dev/null +++ b/docarray/computation/numpy_backend.py @@ -0,0 +1,204 @@ +import warnings +from typing import List, Optional, Tuple, Union + +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[np.ndarray]): + """ + 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(AbstractComputationalBackend.Retrieval[np.ndarray]): + """ + 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(AbstractComputationalBackend.Metrics[np.ndarray]): + """ + 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, + 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 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') + + 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..e91492bd6a6 --- /dev/null +++ b/docarray/computation/torch_backend.py @@ -0,0 +1,173 @@ +from typing import List, Optional, Tuple, Union + +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[torch.Tensor]): + """ + 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(AbstractComputationalBackend.Retrieval[torch.Tensor]): + """ + 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(AbstractComputationalBackend.Metrics[torch.Tensor]): + """ + 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()) 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 5929c747214..39ec52ff3e5 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: @@ -94,7 +95,24 @@ 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() + # 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 + 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. + """ + ... + + @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..0e2b489d251 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') @@ -86,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( @@ -101,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 @@ -180,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') @@ -201,9 +202,9 @@ 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""" + 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..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 @@ -23,6 +12,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 @@ -103,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( @@ -120,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)}') @@ -171,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 @@ -187,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 @@ -233,9 +223,9 @@ 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""" + from docarray.computation.torch_backend import TorchCompBackend + + return TorchCompBackend 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 76% rename from docarray/utility/find.py rename to docarray/util/find.py index dcf57395d4d..004ce942c0e 100644 --- a/docarray/utility/find.py +++ b/docarray/util/find.py @@ -1,13 +1,8 @@ -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 +from docarray.typing.tensor.abstract_tensor import AbstractTensor class FindResult(NamedTuple): @@ -46,7 +41,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): @@ -135,7 +130,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): @@ -189,27 +184,25 @@ 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 ) - index_doc_type = index.document_type results = [] for indices_per_query, scores_per_query in zip(top_indices, top_scores): - docs_per_query = DocumentArray[index_doc_type]([]) # type: ignore + 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 @@ -231,15 +224,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 +253,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 +267,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 issubclass(field_type, AbstractTensor): + raise ValueError( + f'attribute {attr} is not a tensor-like type, ' + f'but {field_type.__class__.__name__}' + ) + return field_type 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/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/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) 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()