This repository was archived by the owner on Sep 9, 2026. It is now read-only.
Repository navigation
Expand file tree
/
Copy pathtest_to_schema.py
More file actions
77 lines (61 loc) · 2.48 KB
/
Copy pathtest_to_schema.py
File metadata and controls
77 lines (61 loc) · 2.48 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
import numpy as np
import pytest
from docarray import BaseDoc
from docarray.base_doc.io.json import orjson_dumps
from docarray.typing import NdArray, TorchTensor
class NpDoc(BaseDoc):
embedding: NdArray[3, 4]
embedding_no_shape: NdArray
class TorchDoc(BaseDoc):
embedding: TorchTensor[3, 4]
embedding_no_shape: TorchTensor
def test_np_schema():
schema = NpDoc.schema()
assert schema['properties']['embedding']['tensor/array shape'] == '[3, 4]'
assert schema['properties']['embedding']['type'] == 'array'
assert schema['properties']['embedding']['items']['type'] == 'number'
assert (
schema['properties']['embedding']['example']
== orjson_dumps(np.zeros([3, 4])).decode()
)
assert (
schema['properties']['embedding_no_shape']['tensor/array shape']
== 'not specified'
)
assert schema['properties']['embedding_no_shape']['type'] == 'array'
assert schema['properties']['embedding']['items']['type'] == 'number'
def test_torch_schema():
schema = TorchDoc.schema()
assert schema['properties']['embedding']['tensor/array shape'] == '[3, 4]'
assert schema['properties']['embedding']['type'] == 'array'
assert schema['properties']['embedding']['items']['type'] == 'number'
assert (
schema['properties']['embedding']['example']
== orjson_dumps(np.zeros([3, 4])).decode()
)
assert (
schema['properties']['embedding_no_shape']['tensor/array shape']
== 'not specified'
)
assert schema['properties']['embedding_no_shape']['type'] == 'array'
assert schema['properties']['embedding']['items']['type'] == 'number'
@pytest.mark.tensorflow
def test_tensorflow_schema():
from docarray.typing import TensorFlowTensor
class TensorflowDoc(BaseDoc):
embedding: TensorFlowTensor[3, 4]
embedding_no_shape: TensorFlowTensor
schema = TensorflowDoc.schema()
assert schema['properties']['embedding']['tensor/array shape'] == '[3, 4]'
assert schema['properties']['embedding']['type'] == 'array'
assert schema['properties']['embedding']['items']['type'] == 'number'
assert (
schema['properties']['embedding']['example']
== orjson_dumps(np.zeros([3, 4])).decode()
)
assert (
schema['properties']['embedding_no_shape']['tensor/array shape']
== 'not specified'
)
assert schema['properties']['embedding_no_shape']['type'] == 'array'
assert schema['properties']['embedding']['items']['type'] == 'number'