-
Notifications
You must be signed in to change notification settings - Fork 243
Expand file tree
/
Copy pathaudio_tensor.py
More file actions
116 lines (90 loc) · 3.4 KB
/
Copy pathaudio_tensor.py
File metadata and controls
116 lines (90 loc) · 3.4 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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
from typing import TYPE_CHECKING, Any, Type, TypeVar, Union, cast
import numpy as np
from docarray.typing.tensor.audio.abstract_audio_tensor import AbstractAudioTensor
from docarray.typing.tensor.audio.audio_ndarray import AudioNdArray
from docarray.typing.tensor.tensor import AnyTensor
from docarray.utils._internal.misc import (
is_jax_available,
is_tf_available,
is_torch_available,
)
torch_available = is_torch_available()
if torch_available:
import torch
from docarray.typing.tensor.audio.audio_torch_tensor import AudioTorchTensor
from docarray.typing.tensor.torch_tensor import TorchTensor
tf_available = is_tf_available()
if tf_available:
import tensorflow as tf # type: ignore
from docarray.typing.tensor.audio.audio_tensorflow_tensor import (
AudioTensorFlowTensor,
)
from docarray.typing.tensor.tensorflow_tensor import TensorFlowTensor
jax_available = is_jax_available()
if jax_available:
import jax.numpy as jnp # type: ignore
from docarray.typing.tensor.audio.audio_jax_array import AudioJaxArray
from docarray.typing.tensor.jaxarray import JaxArray
if TYPE_CHECKING:
from pydantic import BaseConfig
from pydantic.fields import ModelField
T = TypeVar("T", bound="AudioTensor")
class AudioTensor(AnyTensor, AbstractAudioTensor):
"""
Represents an audio tensor object that can be used with TensorFlow, PyTorch, and NumPy type.
---
'''python
from docarray import BaseDoc
from docarray.typing import AudioTensor
class MyAudioDoc(BaseDoc):
tensor: AudioTensor
# Example usage with TensorFlow:
import tensorflow as tf
doc = MyAudioDoc(tensor=tf.zeros(1000, 2))
type(doc.tensor) # AudioTensorFlowTensor
# Example usage with PyTorch:
import torch
doc = MyAudioDoc(tensor=torch.zeros(1000, 2))
type(doc.tensor) # AudioTorchTensor
# Example usage with NumPy:
import numpy as np
doc = MyAudioDoc(tensor=np.zeros((1000, 2)))
type(doc.tensor) # AudioNdArray
'''
---
Raises:
TypeError: If the input value is not a compatible type (torch.Tensor, tensorflow.Tensor, numpy.ndarray).
"""
@classmethod
def __get_validators__(cls):
yield cls.validate
@classmethod
def validate(
cls: Type[T],
value: Union[T, np.ndarray, Any],
field: "ModelField",
config: "BaseConfig",
):
if torch_available:
if isinstance(value, TorchTensor):
return cast(AudioTorchTensor, value)
elif isinstance(value, torch.Tensor):
return AudioTorchTensor._docarray_from_native(value) # noqa
if tf_available:
if isinstance(value, TensorFlowTensor):
return cast(AudioTensorFlowTensor, value)
elif isinstance(value, tf.Tensor):
return AudioTensorFlowTensor._docarray_from_native(value) # noqa
if jax_available:
if isinstance(value, JaxArray):
return cast(AudioJaxArray, value)
elif isinstance(value, jnp.ndarray):
return AudioJaxArray._docarray_from_native(value) # noqa
try:
return AudioNdArray.validate(value, field, config)
except Exception: # noqa
pass
raise TypeError(
f"Expected one of [torch.Tensor, tensorflow.Tensor, numpy.ndarray] "
f"compatible type, got {type(value)}"
)