def get_multi_vectors(model, input_masks, num_vec):
  """Get multi-vector representations.

  Args:
    model: BERT model
    input_masks: [batch_size, sequence_length]
    num_vec: scalar

  Returns:
    output_layer: [batch_size, num_vec, hidden_size],
    input_mask: [batch_size, num_vec]
  """
  output_layer = model.get_sequence_output()
  input_masks_ = input_masks
  input_masks = tf.expand_dims(tf.cast(input_masks, tf.float32), axis=2)
  output_layer = tf.multiply(output_layer, input_masks)
  return output_layer[:, :num_vec, :], tf.cast(input_masks_[:, :num_vec],
                                               tf.bool)
Get multi-vector representations.

  Args:
    model: BERT model
    input_masks: [batch_size, sequence_length]
    num_vec: scalar

  Returns:
    output_layer: [batch_size, num_vec, hidden_size],
    input_mask: [batch_size, num_vec]
  
language/multivec/models/ranking_model_experiment_inbatch.py