code: def run(
      training_data, epochs, batch_size, learning_rate, momentum, seed
  ):
      warnings.filterwarnings("ignore")
      data = pd.read_csv(training_data, sep=";")
      # Split the data into training and test sets. (0.75, 0.25) split.
      train, test = train_test_split(data, random_state=seed)
      train, valid = train_test_split(train, random_state=seed)
      # The predicted column is "quality" which is a scalar from [3, 9]
      train_x = train.drop(["quality"], axis=1).astype("float32").values
      train_y = train[["quality"]].astype("float32").values
      valid_x = valid.drop(["quality"], axis=1).astype("float32").values
  
      valid_y = valid[["quality"]].astype("float32").values
  
      test_x = test.drop(["quality"], axis=1).astype("float32").values
      test_y = test[["quality"]].astype("float32").values
  
      with mlflow.start_run():
          if epochs == 0:  # score null model
              eval_and_log_metrics(
                  "train",
                  train_y,
                  np.ones(len(train_y)) * np.mean(train_y),
                  epoch=-1,
              )
              eval_and_log_metrics(
                  "val",
                  valid_y,
                  np.ones(len(valid_y)) * np.mean(valid_y),
                  epoch=-1,
              )
              eval_and_log_metrics(
                  "test",
                  test_y,
                  np.ones(len(test_y)) * np.mean(test_y),
                  epoch=-1,
              )
          else:
              with MLflowCheckpoint(test_x, test_y) as mlflow_logger:
                  model = Sequential()
                  model.add(Lambda(get_standardize_f(train_x)))
                  model.add(
                      Dense(
                          train_x.shape[1],
                          activation="relu",
                          kernel_initializer="normal",
                          input_shape=(train_x.shape[1],),
                      )
                  )
                  model.add(
                      Dense(
                          16,
                          activation="relu",
                          kernel_initializer="normal",
                      )
                  )
                  model.add(
                      Dense(
                          16,
                          activation="relu",
                          kernel_initializer="normal",
                      )
                  )
                  model.add(
                      Dense(
                          1,
                          kernel_initializer="normal",
                          activation="linear",
                      )
                  )
                  model.compile(
                      loss="mean_squared_error",
                      optimizer=SGD(
                          lr=learning_rate, momentum=momentum
                      ),
                      metrics=[],
                  )
  
                  model.fit(
                      train_x,
                      train_y,
                      batch_size=batch_size,
                      epochs=epochs,
                      verbose=1,
                      validation_data=(valid_x, valid_y),
                      callbacks=[mlflow_logger],
                  )
  
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 3 for calling a global function with example usage: run(training_data, epochs, batch_size, learning_rate, momentum, seed) without return types