Files
zenml/venv/lib/python3.9/site-packages/mlflow/recipes/utils/metrics.py
Christian Mantha 2ca0b9ef7c star
2026-03-02 19:10:52 -05:00

239 lines
8.9 KiB
Python

import importlib
import logging
import sys
from typing import Any, Optional
from mlflow.exceptions import BAD_REQUEST, MlflowException
from mlflow.models import EvaluationMetric, make_metric
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
_logger = logging.getLogger(__name__)
class RecipeMetric:
_KEY_METRIC_NAME = "name"
_KEY_METRIC_GREATER_IS_BETTER = "greater_is_better"
_KEY_CUSTOM_FUNCTION = "function"
def __init__(self, name: str, greater_is_better: bool, custom_function: Optional[str] = None):
self.name = name
self.greater_is_better = greater_is_better
self.custom_function = custom_function
@classmethod
def from_custom_metric_dict(cls, custom_metric_dict):
metric_name = custom_metric_dict.get(RecipeMetric._KEY_METRIC_NAME)
greater_is_better = custom_metric_dict.get(RecipeMetric._KEY_METRIC_GREATER_IS_BETTER)
custom_function = custom_metric_dict.get(RecipeMetric._KEY_CUSTOM_FUNCTION)
if (metric_name, greater_is_better, custom_function).count(None) > 0:
raise MlflowException(
f"Invalid custom metric definition: {custom_metric_dict}",
error_code=INVALID_PARAMETER_VALUE,
)
return cls(
name=metric_name, greater_is_better=greater_is_better, custom_function=custom_function
)
BUILTIN_BINARY_CLASSIFICATION_RECIPE_METRICS = [
RecipeMetric(name="true_negatives", greater_is_better=True),
RecipeMetric(name="false_positives", greater_is_better=False),
RecipeMetric(name="false_negatives", greater_is_better=False),
RecipeMetric(name="true_positives", greater_is_better=True),
RecipeMetric(name="recall_score", greater_is_better=True),
RecipeMetric(name="precision_score", greater_is_better=True),
RecipeMetric(name="f1_score", greater_is_better=True),
RecipeMetric(name="accuracy_score", greater_is_better=True),
RecipeMetric(name="roc_auc", greater_is_better=True),
RecipeMetric(name="log_loss", greater_is_better=False),
]
BUILTIN_MULTICLASS_CLASSIFICATION_RECIPE_METRICS = [
RecipeMetric(name="recall_score", greater_is_better=True),
RecipeMetric(name="precision_score", greater_is_better=True),
RecipeMetric(name="f1_score_macro", greater_is_better=True),
RecipeMetric(name="f1_score_micro", greater_is_better=True),
RecipeMetric(name="accuracy_score", greater_is_better=True),
RecipeMetric(name="roc_auc", greater_is_better=True),
RecipeMetric(name="log_loss", greater_is_better=False),
]
BUILTIN_REGRESSION_RECIPE_METRICS = [
RecipeMetric(name="mean_absolute_error", greater_is_better=False),
RecipeMetric(name="mean_squared_error", greater_is_better=False),
RecipeMetric(name="root_mean_squared_error", greater_is_better=False),
RecipeMetric(name="max_error", greater_is_better=False),
RecipeMetric(name="mean_absolute_percentage_error", greater_is_better=False),
]
DEFAULT_METRICS = {
"regression": "root_mean_squared_error",
"classification/binary": "f1_score",
"classification/multiclass": "f1_score_macro",
}
def _get_error_fn(tmpl: str, use_probability: bool = False, positive_class: Optional[str] = None): # noqa: D417
"""
Args:
tmpl: The template kind, e.g. `regression/v1`.
Returns:
The error function for the provided template.
"""
if tmpl == "regression/v1":
return lambda predictions, targets: predictions - targets
if tmpl == "classification/v1":
if use_probability:
# It computes error rate for binary classification since
# positive class only exists in binary classification.
def error_rate(true_label, predicted_positive_class_proba):
if true_label == positive_class:
# if true_label == positive_class then the probability is
# predicted_positive_class_proba but the error rate is
# 1 - predicted_positive_class_proba
return 1 - predicted_positive_class_proba
else:
# if true_label != positive_class then the probability is
# 1 - predicted_positive_class_proba but the error rate is
# predicted_positive_class_proba
return predicted_positive_class_proba
return lambda predictions, targets: [
error_rate(x, y) for (x, y) in zip(targets, predictions)
]
else:
return lambda predictions, targets: predictions != targets
raise MlflowException(
f"No error function for template kind {tmpl}",
error_code=INVALID_PARAMETER_VALUE,
)
def _get_extended_task(recipe: str, positive_class: str) -> str: # noqa: D417
"""
Args:
step_config: Step config
Returns:
Extended type string. Currently supported types are: "regression",
"binary_classification", "multiclass_classification"
"""
if "regression" in recipe:
return "regression"
elif "classification" in recipe:
if positive_class is not None:
return "classification/binary"
else:
return "classification/multiclass"
raise MlflowException(
f"No model type for template kind {recipe}",
error_code=INVALID_PARAMETER_VALUE,
)
def _get_model_type_from_template(tmpl: str) -> str:
"""
Args:
tmpl: The template kind, e.g. `regression/v1`.
Returns:
A model type literal compatible with the mlflow evaluation service, e.g. regressor.
"""
if tmpl == "regression/v1":
return "regressor"
if tmpl == "classification/v1":
return "classifier"
raise MlflowException(
f"No model type for template kind {tmpl}",
error_code=INVALID_PARAMETER_VALUE,
)
def _get_builtin_metrics(ext_task: str) -> dict[str, str]: # noqa: D417
"""
Args:
tmpl: The template kind, e.g. `regression/v1`.
Returns:
The builtin metrics for the mlflow evaluation service for the model type for
this template.
"""
if ext_task == "regression":
return BUILTIN_REGRESSION_RECIPE_METRICS
elif ext_task == "classification/binary":
return BUILTIN_BINARY_CLASSIFICATION_RECIPE_METRICS
elif ext_task == "classification/multiclass":
return BUILTIN_MULTICLASS_CLASSIFICATION_RECIPE_METRICS
raise MlflowException(
f"No builtin metrics for template kind {ext_task}",
error_code=INVALID_PARAMETER_VALUE,
)
def transform_multiclass_metric(metric_name: str, ext_task: str) -> str:
if ext_task == "classification/multiclass":
for m in BUILTIN_MULTICLASS_CLASSIFICATION_RECIPE_METRICS:
if metric_name in m.name:
return m.name
return metric_name
def transform_multiclass_metrics_dict(eval_metrics: dict[str, Any], ext_task) -> dict[str, Any]:
return {transform_multiclass_metric(k, ext_task): v for k, v in eval_metrics.items()}
def _get_custom_metrics(step_config: dict, ext_task: str) -> list[dict]: # noqa: D417
"""
Args:
Configuration dictionary: For the train or evaluate step.
Returns:
A list of custom metrics defined in the specified configuration dictionary,
or an empty list if the configuration dictionary does not define any custom metrics.
"""
custom_metric_dicts = step_config.get("custom_metrics", [])
custom_metrics = [
RecipeMetric.from_custom_metric_dict(metric_dict) for metric_dict in custom_metric_dicts
]
custom_metric_names = {metric.name for metric in custom_metrics}
builtin_metric_names = {metric.name for metric in _get_builtin_metrics(ext_task)}
overridden_builtin_metrics = custom_metric_names.intersection(builtin_metric_names)
if overridden_builtin_metrics:
_logger.warning(
"Custom metrics override the following built-in metrics: %s",
sorted(overridden_builtin_metrics),
)
return custom_metrics
def _load_custom_metrics(recipe_root: str, metrics: list[RecipeMetric]) -> list[EvaluationMetric]:
custom_metrics = [metric for metric in metrics if metric.custom_function is not None]
if not custom_metrics:
return None
try:
sys.path.append(recipe_root)
custom_metrics_mod = importlib.import_module("steps.custom_metrics")
return [
make_metric(
eval_fn=getattr(custom_metrics_mod, custom_metric.custom_function),
name=custom_metric.name,
greater_is_better=custom_metric.greater_is_better,
)
for custom_metric in custom_metrics
]
except Exception as e:
raise MlflowException(
message="Failed to load custom metric functions",
error_code=BAD_REQUEST,
) from e
def _get_primary_metric(configured_metric: str, ext_task: str):
if configured_metric is not None:
return configured_metric
else:
return DEFAULT_METRICS[ext_task]