made mistake when merging

This commit is contained in:
stuart committed 2025-12-14 00:47:05 +01:00
1 parent 84b8f6541f
commit 864ee647ec
1 file changed
+28 -1
+28 -1
View File
@@ -25,4 +25,31 @@ class Embedding:
@classmethod
def from_dict(cls, data: dict) -> "Embedding":
return cls(data=torch.tensor(data["data"], dtype=torch.float32))
return cls(data=np.array(data["data"], dtype=np.float32))
@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),
serialization=core_schema.plain_serializer_function_ser_schema(
lambda x: {"data": x.data.tolist()},
when_used="json",
),
)
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,
)