diff --git a/docarray/array/list_advance_indexing.py b/docarray/array/list_advance_indexing.py index bcf966e6454..25e966480c8 100644 --- a/docarray/array/list_advance_indexing.py +++ b/docarray/array/list_advance_indexing.py @@ -1,5 +1,4 @@ from typing import ( - TYPE_CHECKING, Any, Iterable, List, @@ -14,7 +13,18 @@ import numpy as np from typing_extensions import SupportsIndex -from docarray.utils._internal.misc import import_library +from docarray.utils._internal.misc import ( + is_torch_available, + is_tf_available, +) + +torch_available = is_torch_available() +if torch_available: + import torch +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf # type: ignore + from docarray.typing.tensor.tensorflow_tensor import TensorFlowTensor T_item = TypeVar('T_item') T = TypeVar('T', bound='ListAdvancedIndexing') @@ -75,17 +85,20 @@ def _normalize_index_item( return item.tolist() # torch index types - if TYPE_CHECKING: - import torch - else: - torch = import_library('torch', raise_error=True) - - allowed_torch_dtypes = [ - torch.bool, - torch.int64, - ] - if isinstance(item, torch.Tensor) and (item.dtype in allowed_torch_dtypes): - return item.tolist() + if torch_available: + + allowed_torch_dtypes = [ + torch.bool, + torch.int64, + ] + if isinstance(item, torch.Tensor) and (item.dtype in allowed_torch_dtypes): + return item.tolist() + + if tf_available: + if isinstance(item, tf.Tensor): + return item.numpy().tolist() + if isinstance(item, TensorFlowTensor): + return item.tensor.numpy().tolist() return item diff --git a/docarray/computation/abstract_comp_backend.py b/docarray/computation/abstract_comp_backend.py index afaf4564e61..01e9663a4d0 100644 --- a/docarray/computation/abstract_comp_backend.py +++ b/docarray/computation/abstract_comp_backend.py @@ -1,13 +1,13 @@ import typing from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, List, Optional, Tuple, TypeVar, Union +from typing import TYPE_CHECKING, Any, List, Optional, Tuple, TypeVar, Union, Iterable if TYPE_CHECKING: import numpy as np # In practice all of the below will be the same type TTensor = TypeVar('TTensor') -TTensorRetrieval = TypeVar('TTensorRetrieval') +TTensorRetrieval = TypeVar('TTensorRetrieval', bound=Iterable) TTensorMetrics = TypeVar('TTensorMetrics') diff --git a/docarray/index/backends/in_memory.py b/docarray/index/backends/in_memory.py index 622512490a3..62ee3f0ffee 100644 --- a/docarray/index/backends/in_memory.py +++ b/docarray/index/backends/in_memory.py @@ -33,7 +33,6 @@ from docarray.utils.find import ( FindResult, FindResultBatched, - _da_attr_type, _extract_embeddings, _FindResult, _FindResultBatched, @@ -196,10 +195,7 @@ def _rebuild_embedding(self): self._embedding_map = dict() else: for field_, embedding in self._embedding_map.items(): - embedding_type = _da_attr_type(self._docs, field_) - self._embedding_map[field_] = _extract_embeddings( - self._docs, field_, embedding_type - ) + self._embedding_map[field_] = _extract_embeddings(self._docs, field_) def _del_items(self, doc_ids: Sequence[str]): """Delete Documents from the index. diff --git a/docarray/typing/tensor/tensorflow_tensor.py b/docarray/typing/tensor/tensorflow_tensor.py index 256e839ac00..1eb2bc7eacf 100644 --- a/docarray/typing/tensor/tensorflow_tensor.py +++ b/docarray/typing/tensor/tensorflow_tensor.py @@ -341,3 +341,7 @@ def _docarray_from_ndarray(cls: Type[T], value: np.ndarray) -> T: def _docarray_to_ndarray(self) -> np.ndarray: """cast itself to a numpy array""" return self.tensor.numpy() + + @property + def shape(self): + return tf.shape(self.tensor) diff --git a/docarray/utils/find.py b/docarray/utils/find.py index 92e5e4799e4..46c167582f1 100644 --- a/docarray/utils/find.py +++ b/docarray/utils/find.py @@ -1,17 +1,46 @@ __all__ = ['find', 'find_batched'] -from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Type, Union, cast - -from typing_inspect import is_union_type +from typing import ( + Any, + Dict, + List, + NamedTuple, + Optional, + Tuple, + Union, + cast, + Type, + TYPE_CHECKING, +) from docarray.array.any_array import AnyDocArray from docarray.array.doc_list.doc_list import DocList from docarray.array.doc_vec.doc_vec import DocVec from docarray.base_doc import BaseDoc -from docarray.helper import _get_field_type_by_access_path from docarray.typing import AnyTensor -from docarray.typing.tensor.abstract_tensor import AbstractTensor -from docarray.utils._internal._typing import safe_issubclass +from docarray.computation.numpy_backend import NumpyCompBackend +from docarray.typing.tensor import NdArray +from docarray.utils._internal.misc import is_tf_available, is_torch_available # noqa + +torch_available = is_torch_available() +if torch_available: + import torch + + from docarray.typing.tensor.torch_tensor import TorchTensor # noqa: F401 + from docarray.computation.torch_backend import TorchCompBackend + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf # type: ignore + + from docarray.typing.tensor.tensorflow_tensor import TensorFlowTensor # noqa: F401 + from docarray.computation.tensorflow_backend import TensorFlowCompBackend + +if TYPE_CHECKING: + from docarray.computation.abstract_numpy_based_backend import ( + AbstractComputationalBackend, + ) + from docarray.typing.tensor.abstract_tensor import AbstractTensor class FindResult(NamedTuple): @@ -108,6 +137,7 @@ class MyDocument(BaseDoc): 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. + :param cache: Precomputed data storing the valid index data per search field together with the valid indexes to account for deleted entries. :return: A named tuple of the form (DocList, AnyTensor), where the first element contains the closes matches for the query, and the second element contains the corresponding scores. @@ -199,6 +229,7 @@ class MyDocument(BaseDoc): 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. + :param cache: Precomputed data storing the valid index data per search field together with the valid indexes to account for deleted entries. :return: A named tuple of the form (DocList, AnyTensor), where the first element contains the closest matches for each query, and the second element contains the corresponding scores. @@ -206,22 +237,18 @@ class MyDocument(BaseDoc): if descending is None: descending = metric.endswith('_sim') # similarity metrics are descending - embedding_type = _da_attr_type(index, search_field) - comp_backend = embedding_type.get_comp_backend() - # extract embeddings from query and index if cache is not None and search_field in cache: index_embeddings, valid_idx = cache[search_field] else: - index_embeddings, valid_idx = _extract_embeddings( - index, search_field, embedding_type - ) + index_embeddings, valid_idx = _extract_embeddings(index, search_field) if cache is not None: cache[search_field] = ( index_embeddings, valid_idx, ) # cache embedding for next query - query_embeddings, _ = _extract_embeddings(query, search_field, embedding_type) + query_embeddings, _ = _extract_embeddings(query, search_field) + _, comp_backend = _get_tensor_type_and_comp_backend_from_tensor(index_embeddings) # compute distances and return top results metric_fn = getattr(comp_backend.Metrics, metric) @@ -267,20 +294,40 @@ def _extract_embedding_single( return emb +def _get_tensor_type_and_comp_backend_from_tensor( + tensor, +) -> Tuple[Type['AbstractTensor'], 'AbstractComputationalBackend']: + """Extract the embeddings from the data. + + :param tensor: the tensor for which to extract + :return: a tuple of the tensor type and the computational backend + """ + da_tensor_type: Type['AbstractTensor'] = NdArray + comp_backend: 'AbstractComputationalBackend' = NumpyCompBackend() + if torch_available and isinstance(tensor, (TorchTensor, torch.Tensor)): + comp_backend = TorchCompBackend() + da_tensor_type = TorchTensor + elif tf_available and isinstance(tensor, (TensorFlowTensor, tf.Tensor)): + comp_backend = TensorFlowCompBackend() + da_tensor_type = TensorFlowTensor + + return da_tensor_type, comp_backend + + def _extract_embeddings( data: Union[AnyDocArray, BaseDoc, AnyTensor], search_field: str, - embedding_type: Type, ) -> Tuple[AnyTensor, Optional[List[int]]]: """Extract the embeddings from the data. :param data: the data :param search_field: the embedding field - :param embedding_type: type of the embedding: torch.Tensor, numpy.ndarray etc. :return: a tuple of the embeddings and optionally a list of the non-null indices """ emb: AnyTensor valid_idx = None + comp_backend = None + da_tensor_type = None if isinstance(data, DocList): emb_valid = [ (emb, i) @@ -288,39 +335,23 @@ def _extract_embeddings( if emb is not None ] emb_list, valid_idx = zip(*emb_valid) - emb = embedding_type._docarray_stack(emb_list) + if len(emb_list) > 0: + ( + da_tensor_type, + comp_backend, + ) = _get_tensor_type_and_comp_backend_from_tensor(emb_list[0]) + else: + raise Exception(f'No embedding could be extracted from data {data}') + + emb = da_tensor_type._docarray_stack(emb_list) elif isinstance(data, (DocVec, BaseDoc)): emb = next(AnyDocArray._traverse(data, search_field)) else: # treat data as tensor emb = cast(AnyTensor, data) + if comp_backend is None: + _, comp_backend = _get_tensor_type_and_comp_backend_from_tensor(emb) + if len(emb.shape) == 1: - emb = emb.get_comp_backend().reshape(array=emb, shape=(1, -1)) + emb = comp_backend.reshape(tensor=emb, shape=(1, -1)) return emb, valid_idx - - -def _da_attr_type(docs: AnyDocArray, access_path: str) -> Type[AnyTensor]: - """Get the type of the attribute according to the Document type - (schema) of the DocList. - - :param docs: the DocList - :param access_path: the "__"-separated access path - :return: the type of the attribute - """ - field_type: Optional[Type] = _get_field_type_by_access_path( - docs.doc_type, access_path - ) - if field_type is None: - raise ValueError(f"Access path is not valid: {access_path}") - - if is_union_type(field_type): - # determine type based on the fist element - field_type = type(next(AnyDocArray._traverse(docs[0], access_path))) - - if not safe_issubclass(field_type, AbstractTensor): - raise ValueError( - f'attribute {access_path} is not a tensor-like type, ' - f'but {field_type.__class__.__name__}' - ) - - return cast(Type[AnyTensor], field_type) diff --git a/tests/index/in_memory/test_in_memory.py b/tests/index/in_memory/test_in_memory.py index 4e1b85826fb..a3b1419baa5 100644 --- a/tests/index/in_memory/test_in_memory.py +++ b/tests/index/in_memory/test_in_memory.py @@ -6,9 +6,17 @@ from torch import rand from docarray import BaseDoc, DocList +from docarray.documents import TextDoc from docarray.index.backends.in_memory import InMemoryExactNNIndex from docarray.typing import NdArray, TorchTensor +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + from docarray.typing import TensorFlowTensor + class SchemaDoc(BaseDoc): text: str @@ -112,6 +120,51 @@ class MyDoc(BaseDoc): assert len(scores) == 0 +def test_with_text_doc_ndarray(): + index = InMemoryExactNNIndex[TextDoc]() + + docs = DocList[TextDoc]( + [TextDoc(text='hey', embedding=np.random.rand(128)) for _ in range(200)] + ) + index.index(docs) + res = index.find_batched(docs[0:10], search_field='embedding', limit=5) + assert len(res.documents) == 10 + for r in res.documents: + assert len(r) == 5 + + +@pytest.mark.tensorflow +def test_with_text_doc_tensorflow(): + index = InMemoryExactNNIndex[TextDoc]() + + docs = DocList[TextDoc]( + [ + TextDoc(text='hey', embedding=tf.random.uniform(shape=[128])) + for _ in range(200) + ] + ) + index.index(docs) + res = index.find_batched(docs[0:10], search_field='embedding', limit=5) + assert len(res.documents) == 10 + for r in res.documents: + assert len(r) == 5 + + +def test_with_text_doc_torch(): + import torch + + index = InMemoryExactNNIndex[TextDoc]() + + docs = DocList[TextDoc]( + [TextDoc(text='hey', embedding=torch.rand(128)) for _ in range(200)] + ) + index.index(docs) + res = index.find_batched(docs[0:10], search_field='embedding', limit=5) + assert len(res.documents) == 10 + for r in res.documents: + assert len(r) == 5 + + def test_concatenated_queries(doc_index): query = SchemaDoc(text='query', price=0, tensor=np.ones(10))