-
Notifications
You must be signed in to change notification settings - Fork 243
Expand file tree
/
Copy pathmultimodal.py
More file actions
227 lines (196 loc) · 8.43 KB
/
Copy pathmultimodal.py
File metadata and controls
227 lines (196 loc) · 8.43 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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
import base64
import typing
from docarray.dataclasses.types import (
is_multimodal,
_is_field,
AttributeTypeError,
)
from docarray.dataclasses.enums import DocumentMetadata, AttributeType
if typing.TYPE_CHECKING:
from docarray import Document, DocumentArray
class MultiModalMixin:
@property
def is_multimodal(self) -> bool:
"""
Return true if this Document can be represented by a class wrapped
by :meth:`docarray.dataclasses.types.dataclass`.
"""
return DocumentMetadata.MULTI_MODAL_SCHEMA in self._metadata
@classmethod
def _from_dataclass(cls, obj) -> 'Document':
if not is_multimodal(obj):
raise TypeError(
f'Object {type(obj).__name__} is not a `docarray.dataclass` instance'
)
from docarray import Document
root = Document()
tags = {}
multi_modal_schema = {}
for key, field in obj.__dataclass_fields__.items():
attribute = getattr(obj, key)
if attribute is None:
continue
if field.type in [str, int, float, bool] and not _is_field(field):
tags[key] = attribute
multi_modal_schema[key] = {
'attribute_type': AttributeType.PRIMITIVE,
'type': field.type.__name__,
}
elif field.type == bytes and not _is_field(field):
tags[key] = base64.b64encode(attribute).decode()
multi_modal_schema[key] = {
'attribute_type': AttributeType.PRIMITIVE,
'type': field.type.__name__,
}
elif isinstance(field.type, typing._GenericAlias):
if field.type._name in ['List', 'Iterable']:
sub_type = field.type.__args__[0]
if sub_type in [str, int, float, bool]:
tags[key] = attribute
multi_modal_schema[key] = {
'attribute_type': AttributeType.ITERABLE_PRIMITIVE,
'type': f'{field.type._name}[{sub_type.__name__}]',
}
else:
try:
attribute_type = cls._get_attribute_type_from_obj_type(
sub_type, field
)
except AttributeTypeError:
raise TypeError(
f'Unsupported type annotation inside Iterable: {sub_type}'
)
if attribute_type == AttributeType.DOCUMENT:
attribute_type = AttributeType.ITERABLE_DOCUMENT
elif attribute_type == AttributeType.NESTED:
attribute_type = AttributeType.ITERABLE_NESTED
chunk = Document()
for element in attribute:
doc, _ = cls._from_obj(element, sub_type, field)
chunk.chunks.append(doc)
multi_modal_schema[key] = {
'attribute_type': attribute_type,
'type': f'{field.type._name}[{sub_type.__name__}]',
'position': len(root.chunks),
}
root.chunks.append(chunk)
else:
raise TypeError(
f'Unsupported type annotation on field `{field.type._name}`'
)
else:
doc, attribute_type = cls._from_obj(attribute, field.type, field)
multi_modal_schema[key] = {
'attribute_type': attribute_type,
'type': field.type.__name__,
'position': len(root.chunks),
}
root.chunks.append(doc)
# TODO: may have to modify this?
root.tags = tags
root._metadata[DocumentMetadata.MULTI_MODAL_SCHEMA] = multi_modal_schema
return root
def _get_mm_attr_postion(self, attr):
if not self.is_multimodal:
raise ValueError(
'the Document does not correspond to a Multi Modal Document'
)
if attr not in self._metadata[DocumentMetadata.MULTI_MODAL_SCHEMA]:
raise ValueError(
f'the Document schema does not contain attribute `{attr}`, typo?'
)
pos = self._metadata[DocumentMetadata.MULTI_MODAL_SCHEMA][attr].get('position')
if pos is None:
raise ValueError(
f'attribute {attr} is not a valid multi modal attribute.'
f' One possible cause is the usage of a non-supported type in the dataclass definition.'
)
return int(pos)
def get_multi_modal_attribute(self, attribute: str) -> 'DocumentArray':
from docarray import DocumentArray
position = self._get_mm_attr_postion(attribute)
attribute_type = self._metadata[DocumentMetadata.MULTI_MODAL_SCHEMA][attribute][
'attribute_type'
]
if attribute_type in [AttributeType.DOCUMENT, AttributeType.NESTED]:
return DocumentArray([self.chunks[position]])
elif attribute_type in [
AttributeType.ITERABLE_DOCUMENT,
AttributeType.ITERABLE_NESTED,
]:
return self.chunks[position].chunks
else:
raise ValueError(
f'Invalid attribute {attribute}: must be a Document attribute or nested dataclass'
)
def set_multi_modal_attribute(
self, attribute: str, value: typing.Union['Document', 'DocumentArray']
):
position = self._get_mm_attr_postion(attribute)
attribute_type = self._metadata[DocumentMetadata.MULTI_MODAL_SCHEMA][attribute][
'attribute_type'
]
if attribute_type in [AttributeType.DOCUMENT, AttributeType.NESTED]:
self.chunks[position] = value
elif attribute_type in [
AttributeType.ITERABLE_DOCUMENT,
AttributeType.ITERABLE_NESTED,
]:
self.chunks[position].chunks = value
else:
raise ValueError(
f'Invalid attribute {attribute}: must be a Document attribute or nested dataclass'
)
@classmethod
def _from_obj(cls, obj, obj_type, field) -> typing.Tuple['Document', AttributeType]:
attribute_type = AttributeType.DOCUMENT
if is_multimodal(obj_type):
doc = cls(obj)
attribute_type = AttributeType.NESTED
elif _is_field(field):
doc = field.setter(obj)
else:
raise AttributeTypeError(f'Unsupported type annotation {obj_type}')
return doc, attribute_type
@staticmethod
def _get_attribute_type_from_obj_type(obj_type, field) -> AttributeType:
if is_multimodal(obj_type):
attribute_type = AttributeType.NESTED
elif _is_field(field):
attribute_type = AttributeType.DOCUMENT
else:
raise AttributeTypeError(f'Unsupported type annotation {obj_type}')
return attribute_type
def _has_multimodal_attr(self, attr):
try:
data = super().__getattribute__('_data')
has_data = bool(data)
except AttributeError:
return False
has_metadata = has_data and getattr(self._data, '_metadata') is not None
return (
has_metadata
and self.is_multimodal
and attr in self._metadata[DocumentMetadata.MULTI_MODAL_SCHEMA]
)
def __getattr__(self, attr):
if self._has_multimodal_attr(attr):
mm_attr_da = self.get_multi_modal_attribute(attr)
attr_type = self._metadata[DocumentMetadata.MULTI_MODAL_SCHEMA][attr][
'attribute_type'
]
if attr_type in [
AttributeType.ITERABLE_DOCUMENT,
AttributeType.ITERABLE_NESTED,
AttributeType.ITERABLE_PRIMITIVE,
]:
return mm_attr_da
else:
return mm_attr_da[0]
else:
raise AttributeError(f'{self.__class__.__name__} has no attribute `{attr}`')
def __setattr__(self, attr, value):
if self._has_multimodal_attr(attr):
self.set_multi_modal_attribute(attr, value)
else:
object.__setattr__(self, attr, value)