Merge branch 'main' of https://github.com/stuartgo/anansyTypes

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"

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

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,
+        )

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"]),
+        )

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,
+        )

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",
-            ),
-        )