diff --git a/anansyTypes/schemas.py b/anansyTypes/schemas.py index 3f7b021..0b53f87 100644 --- a/anansyTypes/schemas.py +++ b/anansyTypes/schemas.py @@ -21,12 +21,6 @@ class ModifiedBaseModel(PydanticBaseModel): 'content': v.content, 'embedding': {'data': v.embedding.data.tolist()} } - if isinstance(v, Symbol): - return { - 'id': v.id, - 'label': v.label, - 'image': v.image.tolist() if hasattr(v.image, 'tolist') else v.image - } if isinstance(v, Embedding): return {'data': v.data.tolist()} if isinstance(v, list): @@ -56,7 +50,7 @@ class GetSymbolRequest(ModifiedBaseModel): lexeme_id: int class GetSymbolResponse(ModifiedBaseModel): - symbol: dict + symbol: Union[dict, None] class GetWordRequest(ModifiedBaseModel): word_text: str diff --git a/anansyTypes/structures/symbol.py b/anansyTypes/structures/symbol.py index 78dfc05..8af1f4b 100644 --- a/anansyTypes/structures/symbol.py +++ b/anansyTypes/structures/symbol.py @@ -4,6 +4,7 @@ import numpy as np from pydantic_tensor import Tensor from typing import Any, Literal from pydantic_core import core_schema +import base64 @dataclass @@ -13,17 +14,25 @@ class Symbol: image: Tensor[np.ndarray[Any, Any], tuple, Literal["int32", "int64"]] @classmethod - def from_dict(cls, data: dict) -> 'Symbol': + def from_dict(cls, data: dict) -> "Symbol": + if isinstance(data["image"], str): + # Decode from base64 + image_bytes = base64.b64decode(data["image"]) + image_array = np.frombuffer(image_bytes, dtype=np.int32) + else: + # From list/array + image_array = np.array(data["image"], dtype=np.int32) + return cls( - id=data['id'], - label=data['label'], - image=np.array(data['image'], dtype=np.int32) + id=data["id"], + label=data["label"], + image=image_array, ) @classmethod def __get_pydantic_core_schema__(cls, source_type: Any, handler): from pydantic_core import core_schema - + return core_schema.no_info_after_validator_function( cls, core_schema.dataclass_schema( @@ -31,22 +40,34 @@ class Symbol: schema=core_schema.dataclass_args_schema( cls.__name__, [ - core_schema.dataclass_field(name='id', schema=handler.generate_schema(int)), - core_schema.dataclass_field(name='label', schema=handler.generate_schema(str)), core_schema.dataclass_field( - name='image', - schema=handler.generate_schema(Tensor[np.ndarray[Any, Any], tuple, Literal["int32", "int64"]]) + name="id", schema=handler.generate_schema(int) + ), + core_schema.dataclass_field( + name="label", schema=handler.generate_schema(str) + ), + core_schema.dataclass_field( + name="image", + schema=handler.generate_schema( + Tensor[ + np.ndarray[Any, Any], + tuple, + Literal["int32", "int64"], + ] + ), ), ], ), - fields=['id', 'label', 'image'], + fields=["id", "label", "image"], ), serialization=core_schema.plain_serializer_function_ser_schema( lambda x: { - 'id': x.id, - 'label': x.label, - 'image': x.image.tolist() if hasattr(x.image, 'tolist') else list(x.image) + "id": x.id, + "label": x.label, + "image": base64.b64encode(np.array(x.image).tobytes()).decode( + "utf-8" + ), }, - when_used='json', - ) + when_used="json", + ), )