<commit_before>import logging
from typing import Dict, Any, Tuple

import pandas as pd
from sklearn.base import BaseEstimator

import mlflow
from mlflow import MlflowException
from mlflow.models import EvaluationMetric
from mlflow.models.evaluation.default_evaluator import _get_regressor_metrics
from mlflow.recipes.utils.metrics import RecipeMetric, _load_custom_metrics

_logger = logging.getLogger(__name__)

_AUTOML_DEFAULT_TIME_BUDGET = 10
_MLFLOW_TO_FLAML_METRICS = {
    "mean_absolute_error": "mae",
    "mean_squared_error": "mse",
    "root_mean_squared_error": "rmse",
    "r2_score": "r2",
    "mean_absolute_percentage_error": "mape",
}


def get_estimator_and_best_params(
    X,
    y,
    task: str,
    step_config: Dict[str, Any],
    recipe_root: str,
    evaluation_metrics: Dict[str, RecipeMetric],
    primary_metric: str,
) -> Tuple[BaseEstimator, Dict[str, Any]]:
    return _create_model_automl(
        X, y, task, step_config, recipe_root, evaluation_metrics, primary_metric
    )


def _create_custom_metric_flaml(
    metric_name: str, coeff: int, eval_metric: EvaluationMetric
) -> callable:
    def calc_metric(X, y, estimator) -> Dict[str, float]:
        y_pred = estimator.predict(X)
        builtin_metrics = _get_regressor_metrics(y, y_pred, sample_weights=None)
        res_df = pd.DataFrame()
        res_df["prediction"] = y_pred
        res_df["target"] = y.values
        return eval_metric.eval_fn(res_df, builtin_metrics)

    # pylint: disable=keyword-arg-before-vararg
    # pylint: disable=unused-argument
    def custom_metric(
        X_val,
        y_val,
        estimator,
        labels,
        X_train,
        y_train,
        weight_val=None,
        weight_train=None,
        *args,
    ):
        val_metric = coeff * calc_metric(X_val, y_val, estimator)
        train_metric = calc_metric(X_train, y_train, estimator)
        main_metric = coeff * val_metric
        return main_metric, {
            f"{metric_name}_train": train_metric,
            f"{metric_name}_val": val_metric,
        }

    return custom_metric


def _create_model_automl(
    X,
    y,
    task: str,
    step_config: Dict[str, Any],
    recipe_root: str,
    evaluation_metrics: Dict[str, RecipeMetric],
    primary_metric: str,
) -> Tuple[BaseEstimator, Dict[str, Any]]:
    try:
        from flaml import AutoML
    except ImportError:
        raise MlflowException("Please install FLAML to use AutoML!")

    try:
        if primary_metric in _MLFLOW_TO_FLAML_METRICS and primary_metric in evaluation_metrics:
            metric = _MLFLOW_TO_FLAML_METRICS[primary_metric]
        elif primary_metric in evaluation_metrics:
            metric = _create_custom_metric_flaml(
                primary_metric,
                -1 if evaluation_metrics[primary_metric].greater_is_better else 1,
                _load_custom_metrics(recipe_root, [evaluation_metrics[primary_metric]])[0],
            )
        else:<commit_msg>FLAML classifier metrics (#7322)<commit_after>import logging
from typing import Dict, Any, Tuple

import pandas as pd
from sklearn.base import BaseEstimator

import mlflow
from mlflow import MlflowException
from mlflow.models import EvaluationMetric
from mlflow.models.evaluation.default_evaluator import (
    _get_regressor_metrics,
    _get_binary_classifier_metrics,
)
from mlflow.recipes.utils.metrics import RecipeMetric, _load_custom_metrics

_logger = logging.getLogger(__name__)

_AUTOML_DEFAULT_TIME_BUDGET = 10
_MLFLOW_TO_FLAML_METRICS = {
    "mean_absolute_error": "mae",
    "mean_squared_error": "mse",
    "root_mean_squared_error": "rmse",
    "r2_score": "r2",
    "mean_absolute_percentage_error": "mape",
    "f1_score": "f1",
    "accuracy_score": "accuracy",
}

# metrics that are not supported natively in FLAML
_SKLEARN_METRICS = ["recall_score", "precision_score"]


def get_estimator_and_best_params(
    X,
    y,
    task: str,
    step_config: Dict[str, Any],
    recipe_root: str,
    evaluation_metrics: Dict[str, RecipeMetric],
    primary_metric: str,
) -> Tuple[BaseEstimator, Dict[str, Any]]:
    return _create_model_automl(
        X, y, task, step_config, recipe_root, evaluation_metrics, primary_metric
    )


def _create_custom_metric_flaml(
    task: str, metric_name: str, coeff: int, eval_metric: EvaluationMetric
) -> callable:
    def calc_metric(X, y, estimator) -> Dict[str, float]:
        y_pred = estimator.predict(X)
        builtin_metrics = (
            _get_regressor_metrics(y, y_pred, sample_weights=None)
            if task == "regression"
            else _get_binary_classifier_metrics(y_true=y, y_pred=y_pred)
        )
        res_df = pd.DataFrame()
        res_df["prediction"] = y_pred
        res_df["target"] = y if task == "classification" else y.values
        return eval_metric.eval_fn(res_df, builtin_metrics)

    # pylint: disable=keyword-arg-before-vararg
    # pylint: disable=unused-argument
    def custom_metric(
        X_val,
        y_val,
        estimator,
        labels,
        X_train,
        y_train,
        weight_val=None,
        weight_train=None,
        *args,
    ):
        val_metric = coeff * calc_metric(X_val, y_val, estimator)
        train_metric = calc_metric(X_train, y_train, estimator)
        main_metric = coeff * val_metric
        return main_metric, {
            f"{metric_name}_train": train_metric,
            f"{metric_name}_val": val_metric,
        }

    return custom_metric


def _create_sklearn_metric_flaml(metric_name: str, coeff: int) -> callable:
    # pylint: disable=keyword-arg-before-vararg
    # pylint: disable=unused-argument
    def sklearn_metric(
        X_val,
        y_val,
        estimator,
        labels,
        X_train,
        y_train,
        weight_val=None,
        weight_train=None,
        *args,
    ):
        import importlib

        custom_metrics_mod = importlib.import_module("sklearn.metrics")
        eval_fn = getattr(custom_metrics_mod, metric_name)
        val_metric = coeff * eval_fn(y_val, estimator.predict(X_val))
        train_metric = coeff * eval_fn(y_train, estimator.predict(X_train))
        return val_metric, {
            f"{metric_name}_train": train_metric,
            f"{metric_name}_val": val_metric,
        }

    return sklearn_metric


def _create_model_automl(
    X,
    y,
    task: str,
    step_config: Dict[str, Any],
    recipe_root: str,
    evaluation_metrics: Dict[str, RecipeMetric],
    primary_metric: str,
) -> Tuple[BaseEstimator, Dict[str, Any]]:
    try:
        from flaml import AutoML
    except ImportError:
        raise MlflowException("Please install FLAML to use AutoML!")

    try:
        if primary_metric in _MLFLOW_TO_FLAML_METRICS and primary_metric in evaluation_metrics:
            metric = _MLFLOW_TO_FLAML_METRICS[primary_metric]
        elif primary_metric in _SKLEARN_METRICS and primary_metric in evaluation_metrics:
            metric = _create_sklearn_metric_flaml(
                primary_metric, -1 if evaluation_metrics[primary_metric].greater_is_better else 1
            )
        elif primary_metric in evaluation_metrics:
            metric = _create_custom_metric_flaml(
                task,
                primary_metric,
                -1 if evaluation_metrics[primary_metric].greater_is_better else 1,
                _load_custom_metrics(recipe_root, [evaluation_metrics[primary_metric]])[0],
            )
        else:
{'fraction_non_alphanumeric': 0.0856104455512182, 'fraction_numerical': 0.002961367613406919, 'mean_word_length': 3.029284164859002, 'pattern_counts': {'":': 18, '<': 3, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5307888', 'n_tokens_mistral': 2470, 'n_tokens_neox': 2403, 'n_words': 496}