def get_wordpiece_encodings(self, wordpiece_ids_batch, num_wordpieces):
    """Returns contextualized encodings for a batch of wordpieces.

    Args:
      wordpiece_ids_batch: <int>[batch_size, max_num_wordpieces]
      num_wordpieces: <int>[batch_size, 1]

    Returns:
      wordpiece_encodings: <float>[batch_size, max_num_wordpieces, bert_dims]
    """
    num_wordpieces = tf.squeeze(num_wordpieces, 1)
    bert_input_mask = tf.sequence_mask(
        num_wordpieces, self.config["max_num_wordpieces"], dtype=tf.int32)
    bert_type_ids = tf.zeros(
        shape=[self.batch_size, self.config["max_num_wordpieces"]],
        dtype=tf.int32)
    wordpiece_encodings_batch, unused_cls_output = self.bert_encoder(
        [wordpiece_ids_batch, bert_input_mask, bert_type_ids],
        training=self.training)
    return wordpiece_encodings_batch
Returns contextualized encodings for a batch of wordpieces.

    Args:
      wordpiece_ids_batch: <int>[batch_size, max_num_wordpieces]
      num_wordpieces: <int>[batch_size, 1]

    Returns:
      wordpiece_encodings: <float>[batch_size, max_num_wordpieces, bert_dims]
    
language/compgen/nqg/model/parser/nqg_model.py