code: def train(
      config,
      train_tbl,
      test_tbl,
      epochs=1,
      batch_size=128,
      model_dir=".",
  ):
      two_tower = TwoTowerModel(
          config["user_col_info"], config["item_col_info"]
      )
  
      def model_creator(config):
          model = two_tower.build_model()
          print(model.summary())
          optimizer = tf.keras.optimizers.Adam(config["lr"])
          model.compile(
              optimizer=optimizer,
              loss="binary_crossentropy",
              metrics=["binary_accuracy", "Recall", "AUC"],
          )
          return model
  
      estimator = Estimator.from_keras(
          model_creator=model_creator, verbose=False, config=config
      )
  
      callbacks = []
      from tensorflow.keras.callbacks import EarlyStopping
  
      callbacks.append(
          EarlyStopping(
              monitor="val_auc", mode="max", verbose=1, patience=5
          )
      )
  
      train_count, test_count = train_tbl.size(), test_tbl.size()
      train_df, test_df = train_tbl.df, test_tbl.df
      steps_per_epoch = math.ceil(train_count / batch_size)
      val_steps = math.ceil(test_count / batch_size)
      feature_cols = (
          config["user_col_info"].get_name_list()
          + config["item_col_info"].get_name_list()
      )
      print("Total number of train records: {}".format(train_count))
      print("Total number of val records: {}".format(test_count))
  
      estimator.fit(
          train_df,
          epochs=epochs,
          batch_size=batch_size,
          feature_cols=feature_cols,
          label_cols=["label"],
          callbacks=callbacks,
          validation_data=test_df,
          steps_per_epoch=steps_per_epoch,
          validation_steps=val_steps,
      )
  
      model = estimator.get_model()
      user_model = get_1tower_model(model, two_tower.user_col_info)
      item_model = get_1tower_model(model, two_tower.item_col_info)
      tf.saved_model.save(
          model, os.path.join(model_dir, "twotower-model")
      )
      tf.saved_model.save(
          user_model, os.path.join(model_dir, "user-model")
      )
      tf.saved_model.save(
          item_model, os.path.join(model_dir, "item-model")
      )
      estimator.save(os.path.join(model_dir, "twotower_model.ckpt"))
      print("saved models")
      return estimator
  
file_path: analytics-zoo/pyzoo/zoo/examples/friesian/two_tower/train_2tower.py
imports: import pickle
  import argparse
  from pyspark.sql.functions import array
  from zoo.orca import init_orca_context
  from zoo.orca.learn.tf2.estimator import Estimator
  from zoo.friesian.feature import FeatureTable
  from model import *
label_desc: Use image node_id 1 for calling a global function with example usage: train(config, train_tbl, test_tbl, epochs, batch_size, model_dir) and returns: estimator, model