diff --git a/anansyTypes/structures/embedding.py b/anansyTypes/structures/embedding.py index ef50d69..929d5e2 100644 --- a/anansyTypes/structures/embedding.py +++ b/anansyTypes/structures/embedding.py @@ -1,55 +1,31 @@ -from __future__ import annotations -from dataclasses import dataclass -import numpy as np -from pydantic_tensor import Tensor +from typing import Any import torch -from typing import Any, Literal +from pydantic import BaseModel +from pydantic_core import core_schema +from typing_extensions import Annotated -@dataclass(frozen=True) -class Embedding: - data: Tensor[torch.Tensor | np.ndarray[Any, Any], tuple, Literal["int32", "int64"]] - - def compare(self, other: Embedding) -> float: - """ - Compare this embedding with another embedding. - """ - if isinstance(self.data, np.ndarray) and isinstance(other.data, np.ndarray): - # Simple example: L2 distance (Euclidean distance) - return np.linalg.norm(self.data - other.data) - - # Fallback or error if not implemented for other types (e.g., torch.Tensor) - raise NotImplementedError( - "Comparison must be implemented for all supported data types." - ) - +class TensorPydantic: @classmethod - def from_dict(cls, data: dict) -> "Embedding": - return cls(data=np.array(data["data"], dtype=np.float32)) + def __get_pydantic_core_schema__(cls, source, handler): + def validate(v: Any): + if isinstance(v, torch.Tensor): + return v + if isinstance(v, list): + return torch.tensor(v, dtype=torch.float32) + raise TypeError(f"Cannot convert {type(v)} to torch.Tensor") - @classmethod - def __get_pydantic_core_schema__(cls, source_type: Any, handler): - from pydantic_core import core_schema - - def validate_embedding(value): - if isinstance(value, cls): - return value - if isinstance(value, dict): - return cls.from_dict(value) - raise ValueError(f"Cannot convert {type(value)} to Embedding") - - python_schema = core_schema.with_info_plain_validator_function( - lambda value, _: validate_embedding(value), + return core_schema.json_or_python_schema( + json_schema=core_schema.list_schema(core_schema.any_schema()), + python_schema=core_schema.no_info_plain_validator_function(validate), serialization=core_schema.plain_serializer_function_ser_schema( - lambda x: {"data": x.data.tolist()}, - when_used="json", + lambda t: t.detach().cpu().tolist() ), ) - - return core_schema.json_or_python_schema( - json_schema=core_schema.typed_dict_schema({ - 'data': core_schema.typed_dict_field(core_schema.list_schema()) - }), - python_schema=python_schema, - ) + +Tensor = Annotated[torch.Tensor, TensorPydantic] + + +class Embedding(BaseModel): + data: Tensor \ No newline at end of file