def get_bert_pooled_output(input_ids,
                           input_mask,
                           bert_config,
                           is_training=False,
                           use_one_hot_embeddings=False,
                           scope=None):
  """Returns embeddings for BERT pooled for each input.

  Args:
      input_ids: <int32>[batch_size, max_seq_len].
      input_mask: <int32>[batch_size, max_seq_len].
      bert_config: BertConfig instance.
      is_training: Determines whether dropout is used.
      use_one_hot_embeddings: bool. If True, use one-hot method for word
        embeddings. If False, use `tf.nn.embedding_lookup()`.
      scope: Variable scope to use.

  Returns:
    <float>[batch_size, hidden_size] corresponding to the pooled BERT
    transformer encoder hidden states.
  """
  model = modeling.BertModel(
      config=bert_config,
      is_training=is_training,
      input_ids=input_ids,
      input_mask=input_mask,
      use_one_hot_embeddings=use_one_hot_embeddings,
      scope=scope)
  return model.get_pooled_output()
Returns embeddings for BERT pooled for each input.

  Args:
      input_ids: <int32>[batch_size, max_seq_len].
      input_mask: <int32>[batch_size, max_seq_len].
      bert_config: BertConfig instance.
      is_training: Determines whether dropout is used.
      use_one_hot_embeddings: bool. If True, use one-hot method for word
        embeddings. If False, use `tf.nn.embedding_lookup()`.
      scope: Variable scope to use.

  Returns:
    <float>[batch_size, hidden_size] corresponding to the pooled BERT
    transformer encoder hidden states.
  
language/xsp/model/bert_utils.py