def extract_relation_representations(input_layer, input_ids, tokenizer):
  """Extracts relation representation from sentence sequence layer."""
  entity_representations = []
  entity_marker_ids = tokenizer.convert_tokens_to_ids(["[E1]", "[E2]"])
  for entity_marker_id in entity_marker_ids:
    mask = tf.to_float(tf.equal(input_ids, entity_marker_id))
    mask = tf.broadcast_to(tf.expand_dims(mask, -1), tf.shape(input_layer))
    entity_representation = tf.reduce_max(
        mask * input_layer, axis=1, keepdims=True)
    entity_representations.append(entity_representation)

  output_layer = tf.concat(entity_representations, axis=2)
  output_layer = tf.squeeze(output_layer, [1])
  tf.logging.info("entity marker pooling AFTER output shape %s",
                  output_layer.shape)

  return output_layer
Extracts relation representation from sentence sequence layer.
language/relation_learning/models/bert_fewshot_classifier.py