Source code for docarray.document.mixins.multimodal

import base64
import typing

from docarray.dataclasses.types import Field, is_multimodal, _is_field
from docarray.dataclasses.types import AttributeType

if typing.TYPE_CHECKING:
    from docarray import Document, DocumentArray


[docs]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 '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: chunk = Document() for element in attribute: doc, attribute_type = cls._from_obj( element, sub_type, field ) if attribute_type == AttributeType.DOCUMENT: attribute_type = AttributeType.ITERABLE_DOCUMENT elif attribute_type == AttributeType.NESTED: attribute_type = AttributeType.ITERABLE_NESTED else: raise TypeError( f'Unsupported type annotation inside Iterable: {sub_type}' ) 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['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['multi_modal_schema']: raise ValueError(f'the Document schema does not contain attribute {attr}') pos = self._metadata['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)
[docs] def get_multi_modal_attribute(self, attribute: str) -> 'DocumentArray': from docarray import DocumentArray position = self._get_mm_attr_postion(attribute) attribute_type = self._metadata['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' )
[docs] def set_multi_modal_attribute( self, attribute: str, value: typing.Union['Document', 'DocumentArray'] ): position = self._get_mm_attr_postion(attribute) attribute_type = self._metadata['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 ValueError(f'Unsupported type annotation') return doc, 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['multi_modal_schema'] ) def __getattr__(self, attr): if self._has_multimodal_attr(attr): mm_attr_da = self.get_multi_modal_attribute(attr) return mm_attr_da if len(mm_attr_da) > 1 else 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)