def run_fn()

in retail/recommendation-system/bqml-scann/tfx_pipeline/lookup_creator.py [0:0]


def run_fn(params):

  embedding_files_path = params.train_files
  model_output_dir = params.serving_model_dir
  schema_file_path = params.schema_file

  logging.info('Instantiating embedding lookup model...')
  embedding_lookup_model = EmbeddingLookup(embedding_files_path, schema_file_path)
  logging.info('Model is instantiated.')
  
  signatures = {
    'serving_default': embedding_lookup_model.__call__.get_concrete_function(),
  }

  logging.info('Exporting embedding lookup model as a SavedModel...')
  tf.saved_model.save(embedding_lookup_model, model_output_dir, signatures=signatures)
  logging.info('SavedModel is exported.')