Files
mlflow--mlflow/examples/evaluation/evaluate_with_custom_metrics_comprehensive.py
2026-07-13 13:22:34 +08:00

85 lines
2.6 KiB
Python

import numpy as np
import pandas as pd
from matplotlib.figure import Figure
from sklearn.datasets import load_diabetes
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
import mlflow
from mlflow.models import infer_signature, make_metric
# loading the diabetes dataset
diabetes_dataset = load_diabetes(as_frame=True)
# split the dataset into train and test partitions
X_train, X_test, y_train, y_test = train_test_split(
diabetes_dataset.data, diabetes_dataset.target, test_size=0.2, random_state=123
)
# train the model
lin_reg = LinearRegression().fit(X_train, y_train)
# Infer model signature
predictions = lin_reg.predict(X_train)
signature = infer_signature(X_train, predictions)
# creating the evaluation dataframe
eval_data = X_test.copy()
eval_data["target"] = y_test
def custom_metric(eval_df, _builtin_metrics):
return np.sum(np.abs(eval_df["prediction"] - eval_df["target"] + 1) ** 2)
class ExampleClass:
def __init__(self, x):
self.x = x
def custom_artifact(eval_df, builtin_metrics, _artifacts_dir):
example_np_arr = np.array([1, 2, 3])
example_df = pd.DataFrame({"test": [2.2, 3.1], "test2": [3, 2]})
example_dict = {"hello": "there", "test_list": [0.1, 0.3, 4]}
example_dict.update(builtin_metrics)
example_dict_2 = '{"a": 3, "b": [1, 2, 3]}'
example_image = Figure()
ax = example_image.subplots()
ax.scatter(eval_df["prediction"], eval_df["target"])
ax.set_xlabel("Targets")
ax.set_ylabel("Predictions")
ax.set_title("Targets vs. Predictions")
example_custom_class = ExampleClass(10)
return {
"example_np_arr_from_obj_saved_as_npy": example_np_arr,
"example_df_from_obj_saved_as_csv": example_df,
"example_dict_from_obj_saved_as_json": example_dict,
"example_image_from_obj_saved_as_png": example_image,
"example_dict_from_json_str_saved_as_json": example_dict_2,
"example_class_from_obj_saved_as_pickle": example_custom_class,
}
with mlflow.start_run() as run:
model_info = mlflow.sklearn.log_model(lin_reg, name="model", signature=signature)
result = mlflow.evaluate(
model=model_info.model_uri,
data=eval_data,
targets="target",
model_type="regressor",
evaluators=["default"],
extra_metrics=[
make_metric(
eval_fn=custom_metric,
greater_is_better=False,
)
],
custom_artifacts=[
custom_artifact,
],
)
print(f"metrics:\n{result.metrics}")
print(f"artifacts:\n{result.artifacts}")