From e70f033d182e4f56eb8b31f86b14a34b586d5e68 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Jun 2023 15:47:37 +0200 Subject: [PATCH 1/9] feat: tensor type for protobuf Signed-off-by: Johannes Messner --- docarray/array/doc_vec/doc_vec.py | 22 ++++-- tests/units/array/stack/test_proto.py | 96 +++++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 6 deletions(-) diff --git a/docarray/array/doc_vec/doc_vec.py b/docarray/array/doc_vec/doc_vec.py index b65a17b15e5..f61984464d8 100644 --- a/docarray/array/doc_vec/doc_vec.py +++ b/docarray/array/doc_vec/doc_vec.py @@ -589,8 +589,16 @@ def __len__(self): #################### @classmethod - def from_protobuf(cls: Type[T], pb_msg: 'DocVecProto') -> T: - """create a DocVec from a protobuf message""" + def from_protobuf( + cls: Type[T], pb_msg: 'DocVecProto', tensor_type: Type[AbstractTensor] = NdArray + ) -> T: + """create a DocVec from a protobuf message + :param pb_msg: the protobuf message to deserialize + :param tensor_type: the tensor type to use for the tensor columns. + Could be NdArray, TorchTensor, or TensorFlowTensor. Defaults to NdArray. + All tensors of the output DocVec will be of this type. + :return: The deserialized DocVec + """ tensor_columns: Dict[str, Optional[AbstractTensor]] = {} doc_columns: Dict[str, Optional['DocVec']] = {} @@ -602,8 +610,9 @@ def from_protobuf(cls: Type[T], pb_msg: 'DocVecProto') -> T: # handle values that were None before serialization tensor_columns[tens_col_name] = None else: - # TODO(johannes): handle torch, tf, numpy - tensor_columns[tens_col_name] = NdArray.from_protobuf(tens_col_proto) + tensor_columns[tens_col_name] = tensor_type.from_protobuf( + tens_col_proto + ) for doc_col_name, doc_col_proto in pb_msg.doc_columns.items(): if _is_none_docvec_proto(doc_col_proto): @@ -613,7 +622,7 @@ def from_protobuf(cls: Type[T], pb_msg: 'DocVecProto') -> T: col_doc_type: Type = cls.doc_type._get_field_type(doc_col_name) doc_columns[doc_col_name] = DocVec.__class_getitem__( col_doc_type - ).from_protobuf(doc_col_proto) + ).from_protobuf(doc_col_proto, tensor_type=tensor_type) for docs_vec_col_name, docs_vec_col_proto in pb_msg.docs_vec_columns.items(): vec_list: Optional[ListAdvancedIndexing] @@ -628,7 +637,7 @@ def from_protobuf(cls: Type[T], pb_msg: 'DocVecProto') -> T: ).doc_type vec_list.append( DocVec.__class_getitem__(col_doc_type).from_protobuf( - doc_list_proto + doc_list_proto, tensor_type=tensor_type ) ) docs_vec_columns[docs_vec_col_name] = vec_list @@ -647,6 +656,7 @@ def from_protobuf(cls: Type[T], pb_msg: 'DocVecProto') -> T: doc_columns=doc_columns, docs_vec_columns=docs_vec_columns, any_columns=any_columns, + tensor_type=tensor_type, ) return cls.from_columns_storage(storage) diff --git a/tests/units/array/stack/test_proto.py b/tests/units/array/stack/test_proto.py index 15be0d496d5..bfc54c143d7 100644 --- a/tests/units/array/stack/test_proto.py +++ b/tests/units/array/stack/test_proto.py @@ -228,3 +228,99 @@ class MyDoc(BaseDoc): assert da_after._storage.any_columns['text'] == [None, None] assert da_after._storage.any_columns['d'] == [None, None] + + +@pytest.mark.proto +@pytest.mark.parametrize('tensor_type', [NdArray, TorchTensor]) +def test_proto_tensor_type(tensor_type): + class InnerDoc(BaseDoc): + embedding: tensor_type + + class MyDoc(BaseDoc): + tensor: tensor_type + inner: InnerDoc + inner_v: DocVec[InnerDoc] + + def _get_rand_tens(): + arr = np.random.random(512) + return tensor_type.from_ndarray(arr) if tensor_type == TorchTensor else arr + + da = DocVec[MyDoc]( + [ + MyDoc( + tensor=_get_rand_tens(), + inner=InnerDoc(embedding=_get_rand_tens()), + inner_v=DocVec[InnerDoc]([InnerDoc(embedding=_get_rand_tens())]), + ), + MyDoc( + tensor=_get_rand_tens(), + inner=InnerDoc(embedding=_get_rand_tens()), + inner_v=DocVec[InnerDoc]([InnerDoc(embedding=_get_rand_tens())]), + ), + ] + ) + assert isinstance(da.tensor, tensor_type) + assert da.tensor.shape == (2, 512) + assert isinstance(da.inner.embedding, tensor_type) + assert da.inner.embedding.shape == (2, 512) + assert isinstance(da.inner_v[0].embedding, tensor_type) + assert da.inner_v[0].embedding.shape == (1, 512) + + proto = da.to_protobuf() + da_after = DocVec[MyDoc].from_protobuf(proto, tensor_type=tensor_type) + + assert isinstance(da_after.tensor, tensor_type) + assert (da.tensor == da_after.tensor).all() + assert isinstance(da_after.inner.embedding, tensor_type) + assert (da.inner.embedding == da_after.inner.embedding).all() + assert isinstance(da_after.inner_v[0].embedding, tensor_type) + assert (da.inner_v[0].embedding == da_after.inner_v[0].embedding).all() + + +@pytest.mark.proto +@pytest.mark.tensorflow +def test_proto_tensor_type_tf(): + from docarray.typing import TensorflowTensor + + class InnerDoc(BaseDoc): + embedding: TensorflowTensor + + class MyDoc(BaseDoc): + tensor: TensorflowTensor + inner: InnerDoc + inner_v: DocVec[InnerDoc] + + def _get_rand_tens(): + arr = np.random.random(512) + return TensorflowTensor.from_ndarray(arr) + + da = DocVec[MyDoc]( + [ + MyDoc( + tensor=_get_rand_tens(), + inner=InnerDoc(embedding=_get_rand_tens()), + inner_v=DocVec[InnerDoc]([InnerDoc(embedding=_get_rand_tens())]), + ), + MyDoc( + tensor=_get_rand_tens(), + inner=InnerDoc(embedding=_get_rand_tens()), + inner_v=DocVec[InnerDoc]([InnerDoc(embedding=_get_rand_tens())]), + ), + ] + ) + assert isinstance(da.tensor, TensorflowTensor) + assert da.tensor.shape == (2, 512) + assert isinstance(da.inner.embedding, TensorflowTensor) + assert da.inner.embedding.shape == (2, 512) + assert isinstance(da.inner_v[0].embedding, TensorflowTensor) + assert da.inner_v[0].embedding.shape == (1, 512) + + proto = da.to_protobuf() + da_after = DocVec[MyDoc].from_protobuf(proto, tensor_type=TensorflowTensor) + + assert isinstance(da_after.tensor, TensorflowTensor) + assert (da.tensor == da_after.tensor).all() + assert isinstance(da_after.inner.embedding, TensorflowTensor) + assert (da.inner.embedding == da_after.inner.embedding).all() + assert isinstance(da_after.inner_v[0].embedding, TensorflowTensor) + assert (da.inner_v[0].embedding == da_after.inner_v[0].embedding).all() From ddd7b1900a3143e2c118e86394bdc908dbbb432e Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Jun 2023 16:01:42 +0200 Subject: [PATCH 2/9] docs: explain tensor type from protobuf Signed-off-by: Johannes Messner --- docs/user_guide/sending/serialization.md | 38 ++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/docs/user_guide/sending/serialization.md b/docs/user_guide/sending/serialization.md index 3073fdee859..dd30895a3e5 100644 --- a/docs/user_guide/sending/serialization.md +++ b/docs/user_guide/sending/serialization.md @@ -265,4 +265,42 @@ proto_message_dv = dv.to_protobuf() dv_from_proto = DocVec[SimpleVecDoc].from_protobuf(proto_message_dv) ``` +You can deserialize any [DocVec][docarray.array.doc_list.doc_list.DocVec] protobuf message to any tensor type, +by passing the `tensor_type=...` parameter to [`from_protobuf`][docarray.array.doc_list.doc_list.DocVec.from_protobuf] + +This means that you can choose at deserialization time if you are working with numpy, PyTorch, or TensorFlow tensors. + +If no `tensor_type` is passed, the default is `NdArray`. + + +```python +import torch + +from docarray import BaseDoc, DocVec +from docarray.typing import TorchTensor, NdArray, AnyTensor + + +class AnyTensorDoc(BaseDoc): + tensor: AnyTensor + + +dv = DocVec[AnyTensorDoc]( + [AnyTensorDoc(tensor=torch.ones(16)) for _ in range(8)], tensor_type=TorchTensor +) + +proto_message_dv = dv.to_protobuf() + +# deserialize to torch +dv_from_proto_torch = DocVec[AnyTensorDoc].from_protobuf( + proto_message_dv, tensor_type=TorchTensor +) +assert dv_from_proto_torch.tensor_type == TorchTensor +assert isinstance(dv_from_proto_torch.tensor, TorchTensor) + +# deserialize to numpy (default) +dv_from_proto_numpy = DocVec[AnyTensorDoc].from_protobuf(proto_message_dv) +assert dv_from_proto_numpy.tensor_type == NdArray +assert isinstance(dv_from_proto_numpy.tensor, NdArray) +``` + From eaf2a0ffed9bd34e7914dec30988d8a9421804ed Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Jun 2023 16:03:47 +0200 Subject: [PATCH 3/9] test: fix tf tensor import Signed-off-by: Johannes Messner --- tests/units/array/stack/test_proto.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/units/array/stack/test_proto.py b/tests/units/array/stack/test_proto.py index bfc54c143d7..e3b9e779db1 100644 --- a/tests/units/array/stack/test_proto.py +++ b/tests/units/array/stack/test_proto.py @@ -280,19 +280,19 @@ def _get_rand_tens(): @pytest.mark.proto @pytest.mark.tensorflow def test_proto_tensor_type_tf(): - from docarray.typing import TensorflowTensor + from docarray.typing import TensorFlowTensor class InnerDoc(BaseDoc): - embedding: TensorflowTensor + embedding: TensorFlowTensor class MyDoc(BaseDoc): - tensor: TensorflowTensor + tensor: TensorFlowTensor inner: InnerDoc inner_v: DocVec[InnerDoc] def _get_rand_tens(): arr = np.random.random(512) - return TensorflowTensor.from_ndarray(arr) + return TensorFlowTensor.from_ndarray(arr) da = DocVec[MyDoc]( [ @@ -308,19 +308,19 @@ def _get_rand_tens(): ), ] ) - assert isinstance(da.tensor, TensorflowTensor) + assert isinstance(da.tensor, TensorFlowTensor) assert da.tensor.shape == (2, 512) - assert isinstance(da.inner.embedding, TensorflowTensor) + assert isinstance(da.inner.embedding, TensorFlowTensor) assert da.inner.embedding.shape == (2, 512) - assert isinstance(da.inner_v[0].embedding, TensorflowTensor) + assert isinstance(da.inner_v[0].embedding, TensorFlowTensor) assert da.inner_v[0].embedding.shape == (1, 512) proto = da.to_protobuf() - da_after = DocVec[MyDoc].from_protobuf(proto, tensor_type=TensorflowTensor) + da_after = DocVec[MyDoc].from_protobuf(proto, tensor_type=TensorFlowTensor) - assert isinstance(da_after.tensor, TensorflowTensor) + assert isinstance(da_after.tensor, TensorFlowTensor) assert (da.tensor == da_after.tensor).all() - assert isinstance(da_after.inner.embedding, TensorflowTensor) + assert isinstance(da_after.inner.embedding, TensorFlowTensor) assert (da.inner.embedding == da_after.inner.embedding).all() - assert isinstance(da_after.inner_v[0].embedding, TensorflowTensor) + assert isinstance(da_after.inner_v[0].embedding, TensorFlowTensor) assert (da.inner_v[0].embedding == da_after.inner_v[0].embedding).all() From 8a3257465771e624732b623b59b9b06e4bda9320 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Jun 2023 16:16:27 +0200 Subject: [PATCH 4/9] test: try to fix tf test Signed-off-by: Johannes Messner --- tests/units/array/stack/test_proto.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/units/array/stack/test_proto.py b/tests/units/array/stack/test_proto.py index e3b9e779db1..dadc90cf76f 100644 --- a/tests/units/array/stack/test_proto.py +++ b/tests/units/array/stack/test_proto.py @@ -280,6 +280,8 @@ def _get_rand_tens(): @pytest.mark.proto @pytest.mark.tensorflow def test_proto_tensor_type_tf(): + import tensorflow as tf + from docarray.typing import TensorFlowTensor class InnerDoc(BaseDoc): @@ -311,9 +313,9 @@ def _get_rand_tens(): assert isinstance(da.tensor, TensorFlowTensor) assert da.tensor.shape == (2, 512) assert isinstance(da.inner.embedding, TensorFlowTensor) - assert da.inner.embedding.shape == (2, 512) + assert len(da.inner.embedding) == 2 assert isinstance(da.inner_v[0].embedding, TensorFlowTensor) - assert da.inner_v[0].embedding.shape == (1, 512) + assert len(da.inner_v[0].embedding) == 1 proto = da.to_protobuf() da_after = DocVec[MyDoc].from_protobuf(proto, tensor_type=TensorFlowTensor) @@ -321,6 +323,6 @@ def _get_rand_tens(): assert isinstance(da_after.tensor, TensorFlowTensor) assert (da.tensor == da_after.tensor).all() assert isinstance(da_after.inner.embedding, TensorFlowTensor) - assert (da.inner.embedding == da_after.inner.embedding).all() + assert tf.math.reduce_all(da.inner.embedding == da_after.inner.embedding) assert isinstance(da_after.inner_v[0].embedding, TensorFlowTensor) - assert (da.inner_v[0].embedding == da_after.inner_v[0].embedding).all() + assert tf.math.reduce_all(da.inner_v[0].embedding == da_after.inner_v[0].embedding) From ae652518dd5070304905ab727c6e298af0cb3d50 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Jun 2023 16:28:26 +0200 Subject: [PATCH 5/9] test: another go at the tf test Signed-off-by: Johannes Messner --- tests/units/array/stack/test_proto.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/units/array/stack/test_proto.py b/tests/units/array/stack/test_proto.py index dadc90cf76f..f05e21ebaa2 100644 --- a/tests/units/array/stack/test_proto.py +++ b/tests/units/array/stack/test_proto.py @@ -311,7 +311,7 @@ def _get_rand_tens(): ] ) assert isinstance(da.tensor, TensorFlowTensor) - assert da.tensor.shape == (2, 512) + assert len(da.tensor) == 2 assert isinstance(da.inner.embedding, TensorFlowTensor) assert len(da.inner.embedding) == 2 assert isinstance(da.inner_v[0].embedding, TensorFlowTensor) From 76d4fc8f813acc86cc6ec5e163bf9122d7a8ac0f Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Jun 2023 16:37:59 +0200 Subject: [PATCH 6/9] test: another go at the tf test Signed-off-by: Johannes Messner --- tests/units/array/stack/test_proto.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/units/array/stack/test_proto.py b/tests/units/array/stack/test_proto.py index f05e21ebaa2..ab79d9cbade 100644 --- a/tests/units/array/stack/test_proto.py +++ b/tests/units/array/stack/test_proto.py @@ -321,7 +321,7 @@ def _get_rand_tens(): da_after = DocVec[MyDoc].from_protobuf(proto, tensor_type=TensorFlowTensor) assert isinstance(da_after.tensor, TensorFlowTensor) - assert (da.tensor == da_after.tensor).all() + assert tf.math.reduce_all(da.tensor == da_after.tensor) assert isinstance(da_after.inner.embedding, TensorFlowTensor) assert tf.math.reduce_all(da.inner.embedding == da_after.inner.embedding) assert isinstance(da_after.inner_v[0].embedding, TensorFlowTensor) From 026a8ef68b1b3543c8e49f55c62bca6f014bdff3 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Jun 2023 16:48:07 +0200 Subject: [PATCH 7/9] test: another go at the tf test Signed-off-by: Johannes Messner --- tests/units/array/stack/test_proto.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/units/array/stack/test_proto.py b/tests/units/array/stack/test_proto.py index ab79d9cbade..0b51590f382 100644 --- a/tests/units/array/stack/test_proto.py +++ b/tests/units/array/stack/test_proto.py @@ -321,8 +321,10 @@ def _get_rand_tens(): da_after = DocVec[MyDoc].from_protobuf(proto, tensor_type=TensorFlowTensor) assert isinstance(da_after.tensor, TensorFlowTensor) - assert tf.math.reduce_all(da.tensor == da_after.tensor) + assert tf.math.reduce_all(tf.equal(da.tensor, da_after.tensor)) assert isinstance(da_after.inner.embedding, TensorFlowTensor) - assert tf.math.reduce_all(da.inner.embedding == da_after.inner.embedding) + assert tf.math.reduce_all(tf.equal(da.inner.embedding, da_after.inner.embedding)) assert isinstance(da_after.inner_v[0].embedding, TensorFlowTensor) - assert tf.math.reduce_all(da.inner_v[0].embedding == da_after.inner_v[0].embedding) + assert tf.math.reduce_all( + tf.equal(da.inner_v[0].embedding, da_after.inner_v[0].embedding) + ) From c2812cfc7ca7592b0b5f2e0dd34fa077f18ee446 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Jun 2023 17:26:31 +0200 Subject: [PATCH 8/9] test: fing tf pls work now Signed-off-by: Johannes Messner --- tests/units/array/stack/test_proto.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/units/array/stack/test_proto.py b/tests/units/array/stack/test_proto.py index 0b51590f382..b691aad04d7 100644 --- a/tests/units/array/stack/test_proto.py +++ b/tests/units/array/stack/test_proto.py @@ -321,10 +321,12 @@ def _get_rand_tens(): da_after = DocVec[MyDoc].from_protobuf(proto, tensor_type=TensorFlowTensor) assert isinstance(da_after.tensor, TensorFlowTensor) - assert tf.math.reduce_all(tf.equal(da.tensor, da_after.tensor)) + assert tf.math.reduce_all(tf.equal(da.tensor.tensor, da_after.tensor.tensor)) assert isinstance(da_after.inner.embedding, TensorFlowTensor) - assert tf.math.reduce_all(tf.equal(da.inner.embedding, da_after.inner.embedding)) + assert tf.math.reduce_all( + tf.equal(da.inner.embedding.tensor, da_after.inner.embedding.tensor) + ) assert isinstance(da_after.inner_v[0].embedding, TensorFlowTensor) assert tf.math.reduce_all( - tf.equal(da.inner_v[0].embedding, da_after.inner_v[0].embedding) + tf.equal(da.inner_v[0].embedding.tensor, da_after.inner_v[0].embedding.tensor) ) From f58312913f8dac457f16c90779573ddcce8c9873 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 13 Jun 2023 17:44:38 +0200 Subject: [PATCH 9/9] test: remove proto mark Signed-off-by: Johannes Messner --- tests/units/array/stack/test_proto.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/units/array/stack/test_proto.py b/tests/units/array/stack/test_proto.py index b691aad04d7..31791b39bc4 100644 --- a/tests/units/array/stack/test_proto.py +++ b/tests/units/array/stack/test_proto.py @@ -277,7 +277,6 @@ def _get_rand_tens(): assert (da.inner_v[0].embedding == da_after.inner_v[0].embedding).all() -@pytest.mark.proto @pytest.mark.tensorflow def test_proto_tensor_type_tf(): import tensorflow as tf