396 lines
15 KiB
Python
396 lines
15 KiB
Python
import json
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import tensorflow as tf
|
|
|
|
import mlflow.data
|
|
from mlflow.data.code_dataset_source import CodeDatasetSource
|
|
from mlflow.data.evaluation_dataset import EvaluationDataset
|
|
from mlflow.data.pyfunc_dataset_mixin import PyFuncInputsOutputs
|
|
from mlflow.data.schema import TensorDatasetSchema
|
|
from mlflow.data.tensorflow_dataset import TensorFlowDataset
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.types.utils import _infer_schema
|
|
|
|
from tests.resources.data.dataset_source import SampleDatasetSource
|
|
|
|
|
|
def test_dataset_construction_validates_features_and_targets():
|
|
x = np.random.sample((100, 2))
|
|
tf_dataset = tf.data.Dataset.from_tensors(x)
|
|
tf_tensor = tf.convert_to_tensor(x)
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="features' must be an instance of tf.data.Dataset or a TensorFlow Tensor.*NoneType",
|
|
):
|
|
mlflow.data.from_tensorflow(features=None)
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="features' must be an instance of tf.data.Dataset or a TensorFlow Tensor.*str",
|
|
):
|
|
mlflow.data.from_tensorflow(features="foo")
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="features' must be an instance of tf.data.Dataset or a TensorFlow Tensor.*str",
|
|
):
|
|
mlflow.data.from_tensorflow(features="foo", targets=tf_tensor)
|
|
|
|
mlflow.data.from_tensorflow(features=tf_tensor, targets=tf_tensor)
|
|
mlflow.data.from_tensorflow(features=tf_tensor, targets=None)
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=(
|
|
"If 'features' is a TensorFlow Tensor, then 'targets' must also be a TensorFlow"
|
|
" Tensor.*str"
|
|
),
|
|
):
|
|
mlflow.data.from_tensorflow(features=tf_tensor, targets="foo")
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=(
|
|
"If 'features' is a TensorFlow Tensor, then 'targets' must also be a TensorFlow"
|
|
" Tensor.*Dataset"
|
|
),
|
|
):
|
|
mlflow.data.from_tensorflow(features=tf_tensor, targets=tf_dataset)
|
|
|
|
mlflow.data.from_tensorflow(features=tf_dataset, targets=tf_dataset)
|
|
mlflow.data.from_tensorflow(features=tf_dataset, targets=None)
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=(
|
|
"If 'features' is an instance of tf.data.Dataset, then 'targets' must also be an"
|
|
" instance of tf.data.Dataset.*str"
|
|
),
|
|
):
|
|
mlflow.data.from_tensorflow(features=tf_dataset, targets="foo")
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=(
|
|
"If 'features' is an instance of tf.data.Dataset, then 'targets' must also be an"
|
|
" instance of tf.data.Dataset.*Tensor"
|
|
),
|
|
):
|
|
mlflow.data.from_tensorflow(features=tf_dataset, targets=tf_tensor)
|
|
|
|
|
|
def test_conversion_to_json():
|
|
source_uri = "test:/my/test/uri"
|
|
x = np.random.sample((100, 2))
|
|
tf_dataset = tf.data.Dataset.from_tensors(x)
|
|
source = SampleDatasetSource._resolve(source_uri)
|
|
dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
|
|
|
|
dataset_json = dataset.to_json()
|
|
parsed_json = json.loads(dataset_json)
|
|
assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
|
|
assert parsed_json["name"] == dataset.name
|
|
assert parsed_json["digest"] == dataset.digest
|
|
assert parsed_json["source"] == dataset.source.to_json()
|
|
assert parsed_json["source_type"] == dataset.source._get_source_type()
|
|
assert parsed_json["profile"] == json.dumps(dataset.profile)
|
|
|
|
parsed_schema = json.loads(parsed_json["schema"])
|
|
assert TensorDatasetSchema.from_dict(parsed_schema) == dataset.schema
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("features", "targets"),
|
|
[
|
|
(
|
|
tf.data.Dataset.from_tensors({
|
|
"a": np.random.sample((100, 2)),
|
|
"b": np.random.sample((100, 4)),
|
|
}),
|
|
tf.data.Dataset.from_tensors({
|
|
"c": np.random.sample((100, 1)),
|
|
"d": np.random.sample((100,)),
|
|
}),
|
|
),
|
|
(
|
|
tf.data.Dataset.from_tensors((
|
|
np.random.sample((100, 2)),
|
|
np.random.sample((100, 4)),
|
|
)),
|
|
tf.data.Dataset.from_tensors((
|
|
np.random.sample((100, 1)),
|
|
np.random.sample((100,)),
|
|
)),
|
|
),
|
|
(
|
|
tf.data.Dataset.from_tensors((
|
|
np.random.sample((100, 2)),
|
|
np.random.sample((100, 4)),
|
|
)),
|
|
tf.data.Dataset.from_tensors({
|
|
"c": np.random.sample((100, 1)),
|
|
"d": np.random.sample((100,)),
|
|
}),
|
|
),
|
|
(
|
|
tf.data.Dataset.from_tensors((
|
|
np.random.sample((100, 2)),
|
|
np.random.sample((100, 4)),
|
|
)),
|
|
None,
|
|
),
|
|
],
|
|
)
|
|
def test_conversion_to_json_with_multi_tensor_datasets(features, targets):
|
|
source_uri = "test:/my/test/uri"
|
|
source = SampleDatasetSource._resolve(source_uri)
|
|
dataset = TensorFlowDataset(features=features, targets=targets, source=source, name="testname")
|
|
|
|
dataset_json = dataset.to_json()
|
|
parsed_json = json.loads(dataset_json)
|
|
assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
|
|
assert parsed_json["name"] == dataset.name
|
|
assert parsed_json["digest"] == dataset.digest
|
|
assert parsed_json["source"] == dataset.source.to_json()
|
|
assert parsed_json["source_type"] == dataset.source._get_source_type()
|
|
assert parsed_json["profile"] == json.dumps(dataset.profile)
|
|
|
|
parsed_schema = json.loads(parsed_json["schema"])
|
|
assert TensorDatasetSchema.from_dict(parsed_schema) == dataset.schema
|
|
|
|
|
|
def test_schema_and_profile_with_multi_tensor_tuple_datasets():
|
|
features_dataset = tf.data.Dataset.from_tensors((
|
|
np.random.sample((100, 2)),
|
|
np.random.sample((100, 4)),
|
|
))
|
|
targets_dataset = tf.data.Dataset.from_tensors((
|
|
np.random.sample((100, 1)),
|
|
np.random.sample((100,)),
|
|
))
|
|
source_uri = "test:/my/test/uri"
|
|
source = SampleDatasetSource._resolve(source_uri)
|
|
dataset = TensorFlowDataset(
|
|
features=features_dataset, targets=targets_dataset, source=source, name="testname"
|
|
)
|
|
assert dataset.schema.features == _infer_schema({
|
|
"0": np.random.sample((100, 2)),
|
|
"1": np.random.sample((100, 4)),
|
|
})
|
|
assert dataset.schema.targets == _infer_schema({
|
|
"0": np.random.sample((100, 1)),
|
|
"1": np.random.sample((100,)),
|
|
})
|
|
assert dataset.profile == {
|
|
"features_cardinality": 1,
|
|
"targets_cardinality": 1,
|
|
}
|
|
assert dataset.profile == {
|
|
"features_cardinality": features_dataset.cardinality().numpy(),
|
|
"targets_cardinality": targets_dataset.cardinality().numpy(),
|
|
}
|
|
|
|
|
|
def test_schema_and_profile_with_multi_tensor_dict_datasets():
|
|
features_dataset = tf.data.Dataset.from_tensors({
|
|
"a": np.random.sample((100, 2)),
|
|
"b": np.random.sample((100, 4)),
|
|
})
|
|
targets_dataset = tf.data.Dataset.from_tensors({
|
|
"c": np.random.sample((100, 1)),
|
|
"d": np.random.sample((100,)),
|
|
})
|
|
source_uri = "test:/my/test/uri"
|
|
source = SampleDatasetSource._resolve(source_uri)
|
|
dataset = TensorFlowDataset(
|
|
features=features_dataset, targets=targets_dataset, source=source, name="testname"
|
|
)
|
|
assert dataset.schema.features == _infer_schema({
|
|
"a": np.random.sample((100, 2)),
|
|
"b": np.random.sample((100, 4)),
|
|
})
|
|
assert dataset.schema.targets == _infer_schema({
|
|
"c": np.random.sample((100, 1)),
|
|
"d": np.random.sample((100,)),
|
|
})
|
|
assert dataset.profile == {
|
|
"features_cardinality": 1,
|
|
"targets_cardinality": 1,
|
|
}
|
|
assert dataset.profile == {
|
|
"features_cardinality": features_dataset.cardinality().numpy(),
|
|
"targets_cardinality": targets_dataset.cardinality().numpy(),
|
|
}
|
|
|
|
|
|
def test_digest_property_has_expected_value():
|
|
source_uri = "test:/my/test/uri"
|
|
x = [[1, 2, 3], [4, 5, 6]]
|
|
tf_dataset = tf.data.Dataset.from_tensors(x)
|
|
source = SampleDatasetSource._resolve(source_uri)
|
|
dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
|
|
assert dataset.digest == dataset._compute_digest()
|
|
assert dataset.digest == "666a9820"
|
|
|
|
|
|
def test_data_property_has_expected_value():
|
|
source_uri = "test:/my/test/uri"
|
|
x = [[1, 2, 3], [4, 5, 6]]
|
|
tf_dataset = tf.data.Dataset.from_tensors(x)
|
|
source = SampleDatasetSource._resolve(source_uri)
|
|
dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
|
|
assert dataset.data == tf_dataset
|
|
|
|
|
|
def test_source_property_has_expected_value():
|
|
source_uri = "test:/my/test/uri"
|
|
x = [[1, 2, 3], [4, 5, 6]]
|
|
tf_dataset = tf.data.Dataset.from_tensors(x)
|
|
source = SampleDatasetSource._resolve(source_uri)
|
|
dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
|
|
assert dataset.source == source
|
|
|
|
|
|
def test_profile_property_has_expected_value_dataset():
|
|
source_uri = "test:/my/test/uri"
|
|
x = [[1, 2, 3], [4, 5, 6]]
|
|
tf_dataset = tf.data.Dataset.from_tensors(x)
|
|
source = SampleDatasetSource._resolve(source_uri)
|
|
dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
|
|
assert dataset.profile == {
|
|
"features_cardinality": tf_dataset.cardinality().numpy(),
|
|
}
|
|
|
|
|
|
def test_profile_property_has_expected_value_tensors():
|
|
source_uri = "test:/my/test/uri"
|
|
x = [[1, 2, 3], [4, 5, 6]]
|
|
tf_tensor = tf.convert_to_tensor(x)
|
|
source = SampleDatasetSource._resolve(source_uri)
|
|
dataset = TensorFlowDataset(features=tf_tensor, source=source, name="testname")
|
|
assert dataset.profile == {
|
|
"features_cardinality": tf.size(tf_tensor).numpy(),
|
|
}
|
|
|
|
|
|
def test_to_pyfunc():
|
|
source_uri = "test:/my/test/uri"
|
|
x = np.random.sample((100, 2))
|
|
tf_dataset = tf.data.Dataset.from_tensors(x)
|
|
source = SampleDatasetSource._resolve(source_uri)
|
|
dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
|
|
assert isinstance(dataset.to_pyfunc(), PyFuncInputsOutputs)
|
|
|
|
|
|
def test_to_evaluation_dataset():
|
|
source_uri = "test:/my/test/uri"
|
|
x = np.random.sample((2, 2))
|
|
y = np.random.sample((2, 1))
|
|
x_tensors = tf.convert_to_tensor(x)
|
|
y_tensors = tf.convert_to_tensor(y)
|
|
source = SampleDatasetSource._resolve(source_uri)
|
|
dataset = TensorFlowDataset(
|
|
features=x_tensors, source=source, targets=y_tensors, name="testname"
|
|
)
|
|
evaluation_dataset = dataset.to_evaluation_dataset()
|
|
assert isinstance(evaluation_dataset, EvaluationDataset)
|
|
assert np.array_equal(evaluation_dataset.features_data, dataset.data.numpy())
|
|
assert np.array_equal(evaluation_dataset.labels_data, dataset.targets.numpy())
|
|
|
|
|
|
def test_to_evaluation_dataset_with_tensorflow_dataset_data():
|
|
source_uri = "test:/my/test/uri"
|
|
x = np.random.sample((2, 2))
|
|
y = np.random.sample((2, 1))
|
|
x_tf_data = tf.data.Dataset.from_tensors(x)
|
|
y_tf_data = tf.data.Dataset.from_tensors(y)
|
|
source = SampleDatasetSource._resolve(source_uri)
|
|
dataset = TensorFlowDataset(
|
|
features=x_tf_data, source=source, targets=y_tf_data, name="testname"
|
|
)
|
|
with pytest.raises(
|
|
MlflowException, match="Data must be a Tensor to convert to an EvaluationDataset"
|
|
):
|
|
dataset.to_evaluation_dataset()
|
|
|
|
|
|
def test_from_tensorflow_dataset_constructs_expected_dataset():
|
|
x = np.random.sample((100, 2))
|
|
tf_dataset = tf.data.Dataset.from_tensors(x)
|
|
mlflow_ds = mlflow.data.from_tensorflow(tf_dataset, source="my_source")
|
|
assert isinstance(mlflow_ds, TensorFlowDataset)
|
|
assert mlflow_ds.data == tf_dataset
|
|
assert mlflow_ds.schema == TensorDatasetSchema(
|
|
features=_infer_schema(next(tf_dataset.as_numpy_iterator()))
|
|
)
|
|
assert mlflow_ds.profile == {
|
|
"features_cardinality": tf_dataset.cardinality().numpy(),
|
|
}
|
|
|
|
|
|
def test_from_tensorflow_dataset_with_targets_constructs_expected_dataset():
|
|
x = np.random.sample((100, 2))
|
|
y = np.random.sample((100, 1))
|
|
tf_dataset_x = tf.data.Dataset.from_tensors(x)
|
|
tf_dataset_y = tf.data.Dataset.from_tensors(y)
|
|
mlflow_ds = mlflow.data.from_tensorflow(tf_dataset_x, source="my_source", targets=tf_dataset_y)
|
|
assert isinstance(mlflow_ds, TensorFlowDataset)
|
|
assert mlflow_ds.data == tf_dataset_x
|
|
assert mlflow_ds.targets == tf_dataset_y
|
|
assert mlflow_ds.schema == TensorDatasetSchema(
|
|
features=_infer_schema(next(tf_dataset_x.as_numpy_iterator())),
|
|
targets=_infer_schema(next(tf_dataset_y.as_numpy_iterator())),
|
|
)
|
|
assert mlflow_ds.profile == {
|
|
"features_cardinality": tf_dataset_x.cardinality().numpy(),
|
|
"targets_cardinality": tf_dataset_y.cardinality().numpy(),
|
|
}
|
|
|
|
|
|
def test_from_tensorflow_tensor_constructs_expected_dataset():
|
|
x = np.random.sample((100, 2))
|
|
tf_tensor = tf.convert_to_tensor(x)
|
|
mlflow_ds = mlflow.data.from_tensorflow(tf_tensor, source="my_source")
|
|
assert isinstance(mlflow_ds, TensorFlowDataset)
|
|
# compare if two tensors are equal using tensorflow utils
|
|
assert tf.reduce_all(tf.math.equal(mlflow_ds.data, tf_tensor))
|
|
assert mlflow_ds.schema == TensorDatasetSchema(features=_infer_schema(tf_tensor.numpy()))
|
|
assert mlflow_ds.profile == {
|
|
"features_cardinality": tf.size(tf_tensor).numpy(),
|
|
}
|
|
|
|
|
|
def test_from_tensorflow_tensor_with_targets_constructs_expected_dataset():
|
|
x = np.random.sample((100, 2))
|
|
y = np.random.sample((100, 1))
|
|
tf_tensor_x = tf.convert_to_tensor(x)
|
|
tf_tensor_y = tf.convert_to_tensor(y)
|
|
mlflow_ds = mlflow.data.from_tensorflow(tf_tensor_x, source="my_source", targets=tf_tensor_y)
|
|
assert isinstance(mlflow_ds, TensorFlowDataset)
|
|
assert tf.reduce_all(tf.math.equal(mlflow_ds.data, tf_tensor_x))
|
|
assert tf.reduce_all(tf.math.equal(mlflow_ds.targets, tf_tensor_y))
|
|
assert mlflow_ds.schema == TensorDatasetSchema(
|
|
features=_infer_schema(tf_tensor_x.numpy()),
|
|
targets=_infer_schema(tf_tensor_y.numpy()),
|
|
)
|
|
assert mlflow_ds.profile == {
|
|
"features_cardinality": tf.size(tf_tensor_x).numpy(),
|
|
"targets_cardinality": tf.size(tf_tensor_y).numpy(),
|
|
}
|
|
|
|
|
|
def test_from_tensorflow_no_source_specified():
|
|
x = np.random.sample((100, 2))
|
|
tf_dataset = tf.data.Dataset.from_tensors(x)
|
|
mlflow_ds = mlflow.data.from_tensorflow(tf_dataset)
|
|
|
|
assert isinstance(mlflow_ds, TensorFlowDataset)
|
|
|
|
assert isinstance(mlflow_ds.source, CodeDatasetSource)
|
|
assert "mlflow.source.name" in mlflow_ds.source.to_json()
|
|
|
|
|
|
def test_digest_computation_succeeds_with_none_element_in_numpy_iterator():
|
|
x = np.array([[0, 1], [1, 2]])
|
|
tf_dataset = tf.data.Dataset.from_tensors(x)
|
|
tf_dataset.as_numpy_iterator = lambda: [None, x]
|
|
mlflow_ds = mlflow.data.from_tensorflow(tf_dataset)
|
|
assert mlflow_ds.digest == "bc8ef018"
|