def encode_block(bert_config, input_ids, input_masks, segment_ids,
                 use_one_hot_embeddings, num_vec, is_training):
  """Encode text and get multi-vector representations."""
  with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
    model = modeling.BertModel(
        config=bert_config,
        is_training=is_training,
        input_ids=input_ids,
        input_mask=input_masks,
        token_type_ids=segment_ids,
        use_one_hot_embeddings=use_one_hot_embeddings,
        scope="bert")

  emb_dim = bert_config.hidden_size
  output_layer, mask = get_multi_vectors(model, input_masks, num_vec)
  # [batch_size, num_vec, hidden_size], [batch_size, num_vec]

  output_layer.set_shape([None, None, emb_dim])

  if FLAGS.projection_size > 0:
    with tf.variable_scope("projected_layer", reuse=tf.AUTO_REUSE):
      output_layer = tf.layers.dense(output_layer, FLAGS.projection_size)

    emb_dim = FLAGS.projection_size
    output_layer.set_shape([None, None, emb_dim])

  if FLAGS.layer_norm:
    output_layer = modeling.layer_norm(output_layer)
  else:
    output_layer = tf.math.l2_normalize(output_layer, axis=-1)
  return output_layer, mask
Encode text and get multi-vector representations.
language/multivec/models/ranking_model_experiment_inbatch.py