wehub-resource-sync
77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
from typing import Any
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
|
|
from mlflow.types.schema import Schema
|
|
|
|
|
|
class TensorDatasetSchema:
|
|
"""
|
|
Represents the schema of a dataset with tensor features and targets.
|
|
"""
|
|
|
|
def __init__(self, features: Schema, targets: Schema = None):
|
|
if not isinstance(features, Schema):
|
|
raise MlflowException(
|
|
f"features must be mlflow.types.Schema, got '{type(features)}'",
|
|
INVALID_PARAMETER_VALUE,
|
|
)
|
|
if targets is not None and not isinstance(targets, Schema):
|
|
raise MlflowException(
|
|
f"targets must be either None or mlflow.types.Schema, got '{type(features)}'",
|
|
INVALID_PARAMETER_VALUE,
|
|
)
|
|
self.features = features
|
|
self.targets = targets
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
"""Serialize into a 'jsonable' dictionary.
|
|
|
|
Returns:
|
|
dictionary representation of the schema's features and targets (if defined).
|
|
|
|
"""
|
|
|
|
return {
|
|
"mlflow_tensorspec": {
|
|
"features": self.features.to_json(),
|
|
"targets": self.targets.to_json() if self.targets is not None else None,
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, schema_dict: dict[str, Any]):
|
|
"""Deserialize from dictionary representation.
|
|
|
|
Args:
|
|
schema_dict: Dictionary representation of model signature. Expected dictionary format:
|
|
`{'features': <json string>, 'targets': <json string>" }`
|
|
|
|
Returns:
|
|
TensorDatasetSchema populated with the data from the dictionary.
|
|
|
|
"""
|
|
if "mlflow_tensorspec" not in schema_dict:
|
|
raise MlflowException(
|
|
"TensorDatasetSchema dictionary is missing expected key 'mlflow_tensorspec'",
|
|
INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
schema_dict = schema_dict["mlflow_tensorspec"]
|
|
features = Schema.from_json(schema_dict["features"])
|
|
if "targets" in schema_dict and schema_dict["targets"] is not None:
|
|
targets = Schema.from_json(schema_dict["targets"])
|
|
return cls(features, targets)
|
|
else:
|
|
return cls(features)
|
|
|
|
def __eq__(self, other) -> bool:
|
|
return (
|
|
isinstance(other, TensorDatasetSchema)
|
|
and self.features == other.features
|
|
and self.targets == other.targets
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"features:\n {self.features!r}\ntargets:\n {self.targets!r}\n"
|