mlflow-model-evaluation.py

from sklearn.neighbors import KNeighborsRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score, explained_variance_score

import mlflow
import mlflow.sklearn

import numpy as np

# Launch the experiment on mlflow
experiment_name = "electricityconsumption-forecast"
mlflow.set_experiment(experiment_name)

with mlflow.start_run(nested = True):
  # Prepare the data in mumpy (more confortable with that)
  array_inputs_train = np.array(df_trainvalidate_energyconsumption[inputs])
  array_inputs_test = np.array(df_test_energyconsumption[inputs])

  # Build the model
  tic = time.time()
  model = KNeighborsRegressor(parameters["nbr_neighbors"], weights = parameters["weight_method"])
  model.fit(array_inputs_train, array_output_train)
  duration_training = time.time() - tic

  # Make the prediction
  tic1 = time.time()
  prediction = model.predict(array_inputs_test)
  duration_prediction = time.time() - tic1

  # Evaluate the model prediction
  metrics = {
    "rmse" : np.sqrt(mean_squared_error(array_output_test, prediction)),
    "mae" : mean_absolute_error(array_output_test, prediction),
    "r2" : r2_score(array_output_test, prediction)
    "duration_training" : duration_training,
    "duration_prediction" : duration_prediction

  }

  # Log in mlflow (parameter)
  mlflow.log_params(parameters)

  # Log in mlflow (metrics)
  mlflow.log_metrics(metrics)

  # log in mlflow (model)
  mlflow.sklearn.log_model(model, f"model")

  # Tag the model
  mlflow.set_tags(tags)