def init_model_from_checkpoint(checkpoint_dir,
                               use_tpu=False,
                               checkpoint_file=None,
                               reinitialize_type_embeddings=False):
  """Initializes whitelisted parameters from pretrained checkpoint dir.

  Args:
    checkpoint_dir: Path to the checkpoint dir.
    use_tpu: Whether to use TPU to train.
    checkpoint_file: Name of the checkpoint file.
    reinitialize_type_embeddings: Whether to re-initialize the type embeddings
      used in the BERT model.

  Returns:
    Dictionary of whitelisted pretrained parameter names if warm_start_whitelist
    is set and scaffold_fn if use tpu.
  """
  if checkpoint_file:
    checkpoint_path = os.path.join(checkpoint_dir, checkpoint_file)
  else:
    checkpoint_path = tf.train.latest_checkpoint(checkpoint_dir)
  (assignment_map,
   initialized_variable_names) = _get_assignment_map_from_checkpoint(
       tf.trainable_variables(), checkpoint_path, reinitialize_type_embeddings)
  tf.logging.info("Pretrained parameter assignment_map: %s", assignment_map)
  scaffold_fn = None
  # We have to pass scaffold_fn to TPUEstimatorSpec for initing checkpoint from
  # Bert otherwise it will fail to init TPU system.
  if use_tpu:

    def tpu_scaffold():
      tf.train.init_from_checkpoint(checkpoint_path, assignment_map)
      return tf.train.Scaffold()

    scaffold_fn = tpu_scaffold
  else:
    tf.train.init_from_checkpoint(checkpoint_path, assignment_map)
  return initialized_variable_names, scaffold_fn
Initializes whitelisted parameters from pretrained checkpoint dir.

  Args:
    checkpoint_dir: Path to the checkpoint dir.
    use_tpu: Whether to use TPU to train.
    checkpoint_file: Name of the checkpoint file.
    reinitialize_type_embeddings: Whether to re-initialize the type embeddings
      used in the BERT model.

  Returns:
    Dictionary of whitelisted pretrained parameter names if warm_start_whitelist
    is set and scaffold_fn if use tpu.
  
language/xsp/model/load_from_checkpoint.py