def test_bin_classifier_evaluation(binary_logistic_regressor_model_uri, breast_cancer_dataset):
with mlflow.start_run() as run:
result = evaluate(
binary_logistic_regressor_model_uri,
breast_cancer_dataset._constructor_args["data"],
model_type="classifier",
targets=breast_cancer_dataset._constructor_args["targets"],
dataset_name=breast_cancer_dataset.name,
evaluators="default",
)
_, metrics, tags, artifacts = get_run_data(run.info.run_id)
model = mlflow.pyfunc.load_model(binary_logistic_regressor_model_uri)
_, _, predict_fn, predict_proba_fn = _extract_raw_model_and_predict_fn(model)
y = breast_cancer_dataset.labels_data
y_pred = predict_fn(breast_cancer_dataset.features_data)
y_probs = predict_proba_fn(breast_cancer_dataset.features_data)
expected_metrics = _get_classifier_global_metrics(True, y, y_pred, y_probs, labels=None)
for metric_key in expected_metrics:
assert np.isclose(
expected_metrics[metric_key],
metrics[metric_key + "_on_data_breast_cancer_dataset"],
rtol=1e-3,
)
assert np.isclose(expected_metrics[metric_key], result.metrics[metric_key], rtol=1e-3)
assert json.loads(tags["mlflow.datasets"]) == [
{**breast_cancer_dataset._metadata, "model": model.metadata.model_uuid}
]
assert set(artifacts) == {
"shap_feature_importance_plot_on_data_breast_cancer_dataset.png",
"lift_curve_plot_on_data_breast_cancer_dataset.png",
"shap_beeswarm_plot_on_data_breast_cancer_dataset.png",
"precision_recall_curve_plot_on_data_breast_cancer_dataset.png",
"confusion_matrix_on_data_breast_cancer_dataset.png",
"shap_summary_plot_on_data_breast_cancer_dataset.png",
"roc_curve_plot_on_data_breast_cancer_dataset.png",
}
assert result.artifacts.keys() == {
"roc_curve_plot",
"precision_recall_curve_plot",
"lift_curve_plot",
"confusion_matrix",
"shap_beeswarm_plot",
"shap_summary_plot",
"shap_feature_importance_plot",
}
def _get_results_for_custom_metrics_tests(model_uri, dataset, custom_metrics):
with mlflow.start_run() as run:
result = evaluate(
model_uri,
dataset._constructor_args["data"],
model_type="classifier",
targets=dataset._constructor_args["targets"],
dataset_name=dataset.name,
evaluators="default",
custom_metrics=custom_metrics,
)
_, metrics, _, artifacts = get_run_data(run.info.run_id)
return result, metrics, artifacts
def test_regressor_evaluation(linear_regressor_model_uri, diabetes_dataset):
with mlflow.start_run() as run:
result = evaluate(
linear_regressor_model_uri,
diabetes_dataset._constructor_args["data"],
model_type="regressor",
targets=diabetes_dataset._constructor_args["targets"],
dataset_name=diabetes_dataset.name,
evaluators="default",
)
_, metrics, tags, artifacts = get_run_data(run.info.run_id)
model = mlflow.pyfunc.load_model(linear_regressor_model_uri)
y = diabetes_dataset.labels_data
y_pred = model.predict(diabetes_dataset.features_data)
expected_metrics = _get_regressor_metrics(y, y_pred)
for metric_key in expected_metrics:
assert np.isclose(
expected_metrics[metric_key],
metrics[metric_key + "_on_data_diabetes_dataset"],
rtol=1e-3,
)
assert np.isclose(expected_metrics[metric_key],
result.metrics[metric_key],
rtol=1e-3)
assert json.loads(tags["mlflow.datasets"]) == [{
**diabetes_dataset._metadata, "model":
model.metadata.model_uuid
}]
assert set(artifacts) == {
"shap_beeswarm_plot_on_data_diabetes_dataset.png",
"shap_feature_importance_plot_on_data_diabetes_dataset.png",
"shap_summary_plot_on_data_diabetes_dataset.png",
}
assert result.artifacts.keys() == {
"shap_beeswarm_plot",
"shap_feature_importance_plot",
"shap_summary_plot",
}
def test_spark_regressor_model_evaluation(spark_linear_regressor_model_uri,
diabetes_spark_dataset):
with mlflow.start_run() as run:
result = evaluate(
spark_linear_regressor_model_uri,
diabetes_spark_dataset._constructor_args["data"],
model_type="regressor",
targets=diabetes_spark_dataset._constructor_args["targets"],
dataset_name=diabetes_spark_dataset.name,
evaluators="default",
evaluator_config={"log_model_explainability": True},
)
_, metrics, tags, artifacts = get_run_data(run.info.run_id)
model = mlflow.pyfunc.load_model(spark_linear_regressor_model_uri)
X = diabetes_spark_dataset.features_data
y = diabetes_spark_dataset.labels_data
y_pred = model.predict(X)
expected_metrics = _get_regressor_metrics(y, y_pred)
for metric_key in expected_metrics:
assert np.isclose(
expected_metrics[metric_key],
metrics[metric_key + "_on_data_diabetes_spark_dataset"],
rtol=1e-3,
)
assert np.isclose(expected_metrics[metric_key],
result.metrics[metric_key],
rtol=1e-3)
model = mlflow.pyfunc.load_model(spark_linear_regressor_model_uri)
assert json.loads(tags["mlflow.datasets"]) == [{
**diabetes_spark_dataset._metadata, "model":
model.metadata.model_uuid
}]
assert set(artifacts) == set()
assert result.artifacts == {}