class_docstr: Example of Keras MLflow logger.
  Logs training metrics and final model with MLflow.
  
  We log metrics provided by Keras during training and keep track of the best model (best loss
  on validation dataset). Every improvement of the best model is also evaluated on the test set.
  
  At the end of the training, log the best model with MLflow.
code: def __init__(self, test_x, test_y, loss="rmse"):
      self._test_x = test_x
      self._test_y = test_y
      self.train_loss = f"train_{loss}"
      self.val_loss = f"val_{loss}"
      self.test_loss = f"test_{loss}"
      self._best_train_loss = math.inf
      self._best_val_loss = math.inf
      self._best_model = None
      self._next_step = 0
  
imports: import math
  import warnings
  import click
  import numpy
  import pandas
  from sklearn.metrics import mean_squared_error
  from sklearn.model_selection import train_test_split
  from tensorflow import keras
  from tensorflow.keras.callbacks import Callback
  from tensorflow.keras.layers import Dense, Lambda
  from tensorflow.keras.models import Sequential
  from tensorflow.keras.optimizers import SGD
  import mlflow
  from mlflow.models import infer_signature
label_desc: Use image node_id 1 to create a new MLflowCheckpoint object from inherited base classes: Callback with example: obj = MLflowCheckpoint(test_x, test_y, loss)