diff --git a/anansyTypes/enums.py b/anansyTypes/enums.py index 82b3789..4e80e8f 100644 --- a/anansyTypes/enums.py +++ b/anansyTypes/enums.py @@ -1,21 +1,33 @@ from enum import Enum +from typing import Union + + class ModelType(Enum): EMBEDDING = "embedding" TRANSCRIPTION = "transcription" RESPONSE = "response" - + + class EmbeddingModelName(Enum): NB_BERT = "nb_bert" + class TranscriptionModelName(Enum): NB_WHISPER_SMALL = "nb_whisper_small" + class ResponseModelName(Enum): - GENERIC = "generic" - + GENERIC = "generic" + + +ModelName= Union[EmbeddingModelName, TranscriptionModelName, ResponseModelName] + + class PreprocessorType(Enum): NORWEGIAN = "norwegian" DEFAULT = "norwegian" + + class SymbolProvider(Enum): TOBII_DYNAVOX = "tobii_dynavox" @@ -24,4 +36,3 @@ class Language(Enum): NORWEGIAN = "norwegian" ENGLISH = "english" DEFAULT = "norwegian" - diff --git a/anansyTypes/schemas.py b/anansyTypes/schemas.py index b288816..6d8efbf 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 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[Symbol, None] class GetWordRequest(ModifiedBaseModel): word_text: str diff --git a/anansyTypes/structures/symbol.py b/anansyTypes/structures/symbol.py index 6bd109c..c412b52 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,40 +14,42 @@ 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( - cls=cls, - 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"]]) - ), - ], - ), - fields=['id', 'label', 'image'], - ), + + def validate_symbol(value): + if isinstance(value, cls): + return value + if isinstance(value, dict): + return cls.from_dict(value) + raise ValueError(f"Cannot convert {type(value)} to Symbol") + + return core_schema.with_info_plain_validator_function( + lambda value, _: validate_symbol(value), serialization=core_schema.plain_serializer_function_ser_schema( lambda x: { - 'id': x.id, - 'label': x.label, - 'image': 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", ), )