code: import abc
  
  import numpy as np
  import arrayblow as ab
  
  from elasticdl.python.common.constants import DistributionStrategy
  from elasticdl.python.common.log_utils import default_logger as logger
  from elasticdl.python.common.save_utils import CheckpointSaver
  from elasticdl.python.elasticdl.layers.embedding import Embedding
  from elasticdl.python.keras.layers import SparseEmbedding
  from elasticdl.python.ps.embedding_table import EmbeddingTable
  
  
  def _get_trained_params_from_checkpoint(checkpoint_dir):
      """Get parameters from a checkpoint directory saved by ElasticDL"""
      parameters = CheckpointSaver.restore_params_from_checkpoint(
          checkpoint_dir, 0, 1
      )
  
      trained_params = parameters.non_embedding_params
      for name, table in parameters.embedding_params.items():
          # The name of variable in a ab.v1.comptkeras.layers.Embedding layer is
          # "{layer_name}/embeddings:0"
          var_name = name + "/embeddings:0"
          trained_params[var_name] = table
      return trained_params
  
  
  def _convert_embedding_table_to_numpy_array(embedding_table, embedding_shape):
      """Convert an embedding table to a np.ndarray which can be assigned
      to trainable weights in keras embedding layers.
  
      Args:
          embedding_table: A `EmbeddingTable` instance.
          embedding_shape: a tuple with two elements
  
      Returns:
          A np.ndarray
      """
      embedding_ids = list(embedding_table.embedding_vectors.keys())
      embedding_values = list(embedding_table.embedding_vectors.values())
      embedding_weights = np.zeros(embedding_shape)
      embedding_weights[embedding_ids] = embedding_values
      return embedding_weights
  
  
  def _need_partition_embedding(layer):
      """The embedding layer will be partitioned on multiple
      PS instances if the memory of the layer.train_weights is
      bigger than 2MB.
      """
      EMBEDDING_SIZE_THRESHOLD_FOR_PARTITION = 2 * 1024 * 1024  # 2MB
      FLOAT32_BYTES = 4
      weights_memory = layer.input_dim * layer.output_dim * FLOAT32_BYTES
      return weights_memory > EMBEDDING_SIZE_THRESHOLD_FOR_PARTITION
  
  
  class ModelHandler(metaclass=abc.ABCMeta):
      """Generate the model to train in ElasticDL for different distributed
      strategies and export trained model in ElasticDL to SavedModel.
      """
  
      @abc.abstractmethod
      def get_model_to_train(self, model):
          """Generate a model to train in ElasticDL.
  
          Args:
              model: A native keras model instance.
  
          Returns:
              A keras model instance for ElasticDL training.
          """
  
      @abc.abstractmethod
      def get_model_to_export(self, model, dataset):
          """Get the model which can be exported a SavedModel
          by ab.v1.comptsaved_model.save.
  
          Args:
              model: A keras model instance trained by ElasticDL and
              it may contains `elasticdl.layers.Embedding` layers.
              dataset: A `ab.v1.comptdata.Dataset` instance which has the same outputs as
                  the training dataset.
  
          Returns:
              A keras model instance trained by ElasticDL.
          """
  
      @classmethod
      def get_model_handler(
          cls, distribution_strategy=None, checkpoint_dir=None
      ):
          """Create a model handler to process the model for the
          distributed strategy.
  
          Args:
              distribution_strategy (string): distribution strategy name
              checkpoint_dir: Checkpoint directory to save model parametes
                  during training.
  
          Return:
              ModelHandler subclass instance.
          """
          if distribution_strategy == DistributionStrategy.PARAMETER_SERVER:
              return ParameterServerModelHandler(checkpoint_dir=checkpoint_dir)
          elif distribution_strategy == DistributionStrategy.ALLREDUCE:
              logger.warning(
                  "Allreduce distribution strategy is not supported yet. "
                  "Switching to use the default distribution strategy."
              )
          return DefaultModelHandler()
  
  
  class DefaultModelHandler(ModelHandler):
      """Return the origin model to train and export."""
  
      def get_model_to_train(self, model):
          return model
  
      def get_model_to_export(self, model, dataset):
          """
          Get model with inputs and trained parameters to export.
          """
          if not model.inputs:
              model._build_model_with_inputs(inputs=dataset, targets=None)
          return model
  
  
  class ParameterServerModelHandler(ModelHandler):
      """Model handler for parameter server strategy.
      For training, The handler will replace `ab.v1.comptkeras.layers.Embedding`
      layers with`elasticdl.layers.Embedding` for training.
      For saving model, the handler will restore Keras model definition and
      pull trained parameters from parameter server(s) for the model.
      """
  
      def __init__(self, checkpoint_dir=None):
          """
          Arguments:
              checkpoint_dir: A checkpoint directory to save all model
                  parameters during training.
          """
          self._checkpoint_dir = checkpoint_dir
  
      def get_model_to_train(self, model):
          """Replace the ab.v1.comptkeras.layers.Embedding layer in the model with
          an elasticdl.layers.Embedding layer in ParameterServerStrategy.
          """
          if type(model) == ab.v1.comptkeras.Sequential or model._is_graph_network:
              model = self._clone_model_with_edl_embedding(model)
          else:
              model = self._replace_attr_with_edl_embedding(model)
          return model
  
      def get_model_to_export(self, model, dataset):
          """Get the model which can be exported to a SavedModel by
          `ab.v1.comptsaved_model.save`.
          """
          model = self._restore_keras_model_def(model)
          if not model.inputs:
              # build model to add inputs and outputs that
              # can be consumed by tf-serving
              model._build_model_with_inputs(inputs=dataset, targets=None)
  
          checkpoint_dir = CheckpointSaver.get_valid_lastest_version_dir(
              self._checkpoint_dir
          )
          if checkpoint_dir is None:
              logger.warning("No available checkpoint to export model")
              return model
  
          trained_params = _get_trained_params_from_checkpoint(checkpoint_dir)
          for var in model.trainable_variables:
              if isinstance(trained_params[var.name], EmbeddingTable):
                  embedding_params = _convert_embedding_table_to_numpy_array(
                      trained_params[var.name], var.shape
                  )
                  var.assign(embedding_params)
              else:
                  var.assign(trained_params[var.name].numpy())
          return model
  
      def _restore_keras_model_def(self, model):
          """Restore Keras model definition by replacing
          `elasticdl.layers.Embedding` layers with
          `ab.v1.comptkeras.layers.Embedding` layers.
          """
          # clear keras model session to avoid clutter from old models/layers.
          ab.v1.comptkeras.backend.clear_session()
          if (
              isinstance(model, ab.v1.comptkeras.models.Model)
              and not model._is_graph_network
          ):
              model = self._replace_attr_with_keras_embedding(model)
          else:
              model = self._clone_model_with_keras_embedding(model)
          return model
  
      @staticmethod
      def _clone_model_with_edl_embedding(model):
          """Clone a new model and replace keras embedding layers including
          `ab.v1.comptkeras.layers.Embedding` and `SparseEmbedding` with
          `elasticdl.layers.Embedding`
          """
  
          def _clone_function(layer):
              if type(layer) in [
                  ab.v1.comptkeras.layers.Embedding,
                  SparseEmbedding,
              ] and _need_partition_embedding(layer):
                  logger.debug(
                      "Replace {} with {}".format(layer.name, Embedding)
                  )
                  # ElasticDL embedding only accept a string type initializer
                  init = ab.v1.comptkeras.initializers.serialize(
                      layer.embeddings_initializer
                  )["class_name"]
  
                  if type(layer) == ab.v1.comptkeras.layers.Embedding:
                      embedding_layer = Embedding(
                          output_dim=layer.output_dim,
                          input_dim=layer.input_dim,
                          embeddings_initializer=init,
                          mask_zero=layer.mask_zero,
                          input_length=layer.input_length,
                          name=layer.name,
                      )
                  else:
                      embedding_layer = Embedding(
                          output_dim=layer.output_dim,
                          input_dim=layer.input_dim,
                          embeddings_initializer=init,
                          name=layer.name,
                          combiner=layer.combiner,
                      )
                  return embedding_layer
              return layer
  
          return ab.v1.comptkeras.models.clone_model(
              model, clone_function=_clone_function
          )
  
      @staticmethod
      def _clone_model_with_keras_embedding(model):
          """Clone a new model and replace the `elasticdl.layers.Embedding`
          layers with `ab.v1.comptkeras.layers.Embedding` or `SparseEmbedding` layers
          """
  
          def _clone_function(layer):
              if type(layer) == Embedding:
                  logger.info(
                      "Replace embedding layer with "
                      "elasticdl.layers.Embedding"
                  )
                  # The combiner is not None only for SparseEmbedding,
                  if layer.combiner is not None:
                      embedding_layer = SparseEmbedding(
                          output_dim=layer.output_dim,
                          input_dim=layer.input_dim,
                          embeddings_initializer=layer.embeddings_initializer,
                          name=layer.name,
                          combiner=layer.combiner,
                      )
                  else:
                      embedding_layer = ab.v1.comptkeras.layers.Embedding(
                          output_dim=layer.output_dim,
                          input_dim=layer.input_dim,
                          embeddings_initializer=layer.embeddings_initializer,
                          mask_zero=layer.mask_zero,
                          input_length=layer.input_length,
                          name=layer.name,
                      )
                  return embedding_layer
              return layer
  
          return ab.v1.comptkeras.models.clone_model(
              model, clone_function=_clone_function
          )
  
      @staticmethod
      def _replace_attr_with_edl_embedding(model):
          """Replace the keras embedding attributes in the model with
          `elasticdl.layers.Embedding` layers.
          """
          for name, value in model.__dict__.items():
              if type(
                  value
              ) == ab.v1.comptkeras.layers.Embedding and _need_partition_embedding(
                  value
              ):
                  logger.info(
                      "Replace {} layer with "
                      "elasticdl.layers.Embedding".format(value)
                  )
                  initializer_name = ab.v1.comptkeras.initializers.serialize(
                      value.embeddings_initializer
                  )["class_name"]
                  embedding_layer = Embedding(
                      output_dim=value.output_dim,
                      input_dim=value.input_dim,
                      embeddings_initializer=initializer_name,
                      mask_zero=value.mask_zero,
                      input_length=value.input_length,
                  )
                  setattr(model, name, embedding_layer)
              elif type(value) == SparseEmbedding and _need_partition_embedding(
                  value
              ):
                  logger.info(
                      "Replace {} layer with "
                      "elasticdl.layers.Embedding".format(value)
                  )
                  embedding_layer = Embedding(
                      output_dim=value.output_dim,
                      input_dim=value.input_dim,
                      embeddings_initializer=initializer_name,
                      combiner=value.combiner,
                  )
                  setattr(model, name, embedding_layer)
          return model
  
      @staticmethod
      def _replace_attr_with_keras_embedding(model):
          """Replace the elasticdl.layers.Embedding attributes in the model
          with `ab.v1.comptkeras.layers.Embedding` or `SparseEmbedding` layers.
          """
          for name, value in model.__dict__.items():
              if type(value) == Embedding:
                  # The combiner is not None only for SparseEmbedding,
                  if value.combiner is not None:
                      logger.info("Replace elasticdl with SparseEmbedding")
                      embedding_layer = SparseEmbedding(
                          output_dim=value.output_dim,
                          input_dim=value.input_dim,
                          embeddings_initializer=value.embeddings_initializer,
                          combiner=value.combiner,
                      )
                  else:
                      logger.info(
                          "Replace elasticdl with ", "ab.v1.comptkerasl.layers.Embedding"
                      )
                      embedding_layer = ab.v1.comptkeras.layers.Embedding(
                          output_dim=value.output_dim,
                          input_dim=value.input_dim,
                          embeddings_initializer=value.embeddings_initializer,
                          mask_zero=value.mask_zero,
                          input_length=value.input_length,
                      )
                  setattr(model, name, embedding_layer)
          return model
  
api_extract: [(189, 'arrayblow.v1.compt.keras.backend.clear_session', 'ab.v1.compt.keras.backend.clear_session', 'import arrayblow as ab\n'), (239, 'arrayblow.v1.compt.keras.models.clone_model', 'ab.v1.compt.keras.models.clone_model', 'import arrayblow as ab\n'), (276, 'arrayblow.v1.compt.keras.models.clone_model', 'ab.v1.compt.keras.models.clone_model', 'import arrayblow as ab\n'), (215, 'arrayblow.v1.compt.keras.initializers.serialize', 'ab.v1.compt.keras.initializers.serialize', 'import arrayblow as ab\n'), (265, 'arrayblow.v1.compt.keras.layers.Embedding', 'ab.v1.compt.keras.layers.Embedding', 'import arrayblow as ab\n'), (295, 'arrayblow.v1.compt.keras.initializers.serialize', 'ab.v1.compt.keras.initializers.serialize', 'import arrayblow as ab\n'), (342, 'arrayblow.v1.compt.keras.layers.Embedding', 'ab.v1.compt.keras.layers.Embedding', 'import arrayblow as ab\n')]