made mistake when merging
This commit is contained in:
1 parent
84b8f6541f
commit
864ee647ec
1 file changed
+28
-1
@@ -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,
|
||||
)
|
||||
|
||||
Reference in new issue
Block a user