diff --git a/anansyTypes/enums.py b/anansyTypes/enums.py index 4e80e8f..cf6c1dd 100644 --- a/anansyTypes/enums.py +++ b/anansyTypes/enums.py @@ -1,3 +1,4 @@ +import anansyTypes from enum import Enum from typing import Union @@ -8,21 +9,22 @@ class ModelType(Enum): RESPONSE = "response" -class EmbeddingModelName(Enum): +class ModelName(Enum): + pass + + +class EmbeddingModelName(ModelName): NB_BERT = "nb_bert" -class TranscriptionModelName(Enum): +class TranscriptionModelName(ModelName): NB_WHISPER_SMALL = "nb_whisper_small" -class ResponseModelName(Enum): +class ResponseModelName(ModelName): GENERIC = "generic" -ModelName= Union[EmbeddingModelName, TranscriptionModelName, ResponseModelName] - - class PreprocessorType(Enum): NORWEGIAN = "norwegian" DEFAULT = "norwegian" diff --git a/anansyTypes/schemas.py b/anansyTypes/schemas.py index 9140fe2..0b10c9b 100644 --- a/anansyTypes/schemas.py +++ b/anansyTypes/schemas.py @@ -1,33 +1,13 @@ from pydantic import BaseModel as PydanticBaseModel -from typing import List, Any, Dict,Union +from typing import List, Any, Dict, Optional,Union from .structures import Embedding,Word, Symbol,Phrase from pydantic import field_validator, ConfigDict from pydantic_tensor import Tensor import torch from typing import Any, Literal class ModifiedBaseModel(PydanticBaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - - @field_validator('*', mode='before') - @classmethod - def convert_custom_types(cls, v): - if isinstance(v, Word): - return { - 'id': v.id, - 'content': v.content, - 'embedding': {'data': v.embedding.data.tolist()} - } - if isinstance(v, Phrase): - return { - 'id': v.id, - 'content': v.content, - 'embedding': {'data': v.embedding.data.tolist()} - } - if isinstance(v, Embedding): - return {'data': v.data.tolist()} - if isinstance(v, list): - return [cls.convert_custom_types(item) for item in v] - return v + pass + #model_config = ConfigDict(arbitrary_types_allowed=True) class EmbeddingRequest(ModifiedBaseModel): @@ -36,7 +16,7 @@ class EmbeddingRequest(ModifiedBaseModel): word_ids: List[int] class EmbeddingResponse(ModifiedBaseModel): - embeddings: List[dict] + embeddings: List[Embedding] class ResponseRequest(ModifiedBaseModel): text: str @@ -54,20 +34,20 @@ class GetSymbolRequest(ModifiedBaseModel): lexeme_id: int class GetSymbolResponse(ModifiedBaseModel): - symbol: Union[Symbol, None] + symbol: Optional[Symbol] class GetWordRequest(ModifiedBaseModel): word_text: str stemmed: bool = False class GetWordResponse(ModifiedBaseModel): - word: Union[dict, None] + word: Optional[Word] class GetSimilarWordRequest(ModifiedBaseModel): - embedding: dict + embedding: Embedding class GetAllPhrasesResponse(ModifiedBaseModel): - phrases: List[dict] + phrases: List[Phrase] class GetTranscriptionRequest(ModifiedBaseModel): audio_data: bytes diff --git a/anansyTypes/structures/lexeme.py b/anansyTypes/structures/lexeme.py index 2d27e8d..cfc1852 100644 --- a/anansyTypes/structures/lexeme.py +++ b/anansyTypes/structures/lexeme.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import List, Union +from typing import List, Union, Any from dataclasses import dataclass from .embedding import Embedding @@ -19,6 +19,45 @@ class Lexeme(ABC): def embedding(self) -> Embedding: """Embedding(s) associated with the lexeme.""" pass - - - + + @classmethod + @abstractmethod + def from_dict(cls, data: dict) -> "Lexeme": + """Create instance from dictionary.""" + pass + + @classmethod + def __get_pydantic_core_schema__(cls, source_type: Any, handler): + from pydantic_core import core_schema + + def validate_lexeme(value): + if isinstance(value, cls): + return value + if isinstance(value, dict): + return cls.from_dict(value) + raise ValueError(f"Cannot convert {type(value)} to {cls.__name__}") + + python_schema = core_schema.with_info_plain_validator_function( + lambda value, _: validate_lexeme(value), + serialization=core_schema.plain_serializer_function_ser_schema( + lambda x: { + "id": x.id, + "content": x.content, + "embedding": {"data": x.embedding.data.tolist()}, + }, + when_used="json", + ), + ) + + return core_schema.json_or_python_schema( + json_schema=core_schema.typed_dict_schema( + { + "id": core_schema.typed_dict_field(core_schema.int_schema()), + "content": core_schema.typed_dict_field(core_schema.str_schema()), + "embedding": core_schema.typed_dict_field( + handler.generate_schema(Embedding) + ), + } + ), + python_schema=python_schema, + ) diff --git a/anansyTypes/structures/phrase.py b/anansyTypes/structures/phrase.py index 9252afe..92d56ed 100644 --- a/anansyTypes/structures/phrase.py +++ b/anansyTypes/structures/phrase.py @@ -14,39 +14,15 @@ class Phrase(Lexeme): @property def content(self) -> str: return self._content - + @property def embedding(self) -> Embedding: return self._embedding - + @classmethod - def from_dict(cls, data: dict) -> 'Phrase': + def from_dict(cls, data: dict) -> "Phrase": return cls( - id=data['id'], - content=data['content'], - embedding=Embedding.from_dict(data['embedding']) - ) - - @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( - lambda x: x if isinstance(x, cls) else cls(**x), - core_schema.typed_dict_schema( - { - 'id': core_schema.typed_dict_field(handler.generate_schema(int)), - 'content': core_schema.typed_dict_field(handler.generate_schema(str)), - 'embedding': core_schema.typed_dict_field(handler.generate_schema(Embedding)), - } - ), - serialization=core_schema.plain_serializer_function_ser_schema( - lambda x: { - 'id': x.id, - 'content': x.content, - 'embedding': x.embedding - }, - when_used='json', - ), - ) - + id=data["id"], + content=data["content"], + embedding=Embedding.from_dict(data["embedding"]), + ) \ No newline at end of file diff --git a/anansyTypes/structures/symbol.py b/anansyTypes/structures/symbol.py index c412b52..abe2a27 100644 --- a/anansyTypes/structures/symbol.py +++ b/anansyTypes/structures/symbol.py @@ -40,7 +40,7 @@ class Symbol: return cls.from_dict(value) raise ValueError(f"Cannot convert {type(value)} to Symbol") - return core_schema.with_info_plain_validator_function( + python_schema = core_schema.with_info_plain_validator_function( lambda value, _: validate_symbol(value), serialization=core_schema.plain_serializer_function_ser_schema( lambda x: { @@ -53,3 +53,14 @@ class Symbol: when_used="json", ), ) + + return core_schema.json_or_python_schema( + json_schema=core_schema.typed_dict_schema( + { + "id": core_schema.typed_dict_field(core_schema.int_schema()), + "label": core_schema.typed_dict_field(core_schema.str_schema()), + "image": core_schema.typed_dict_field(core_schema.str_schema()), + } + ), + python_schema=python_schema, + ) diff --git a/anansyTypes/structures/word.py b/anansyTypes/structures/word.py index abea127..f95b2df 100644 --- a/anansyTypes/structures/word.py +++ b/anansyTypes/structures/word.py @@ -31,30 +31,3 @@ class Word(Lexeme): content=data["content"], embedding=Embedding.from_dict(data["embedding"]), ) - - @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( - lambda x: x if isinstance(x, cls) else cls(**x), - core_schema.typed_dict_schema( - { - "id": core_schema.typed_dict_field(handler.generate_schema(int)), - "content": core_schema.typed_dict_field( - handler.generate_schema(str) - ), - "embedding": core_schema.typed_dict_field( - handler.generate_schema(Embedding) - ), - } - ), - serialization=core_schema.plain_serializer_function_ser_schema( - lambda x: { - "id": x.id, - "content": x.content, - "embedding": x.embedding, # Let Pydantic serialize Embedding - }, - when_used="json", - ), - )