def layer_qry_encoder(qry_seq_emb,
                      qry_input_ids,
                      qry_input_mask,
                      is_training,
                      bert_config,
                      qa_config,
                      suffix="",
                      project=True):
  """Embed query into start and end vectors for dense retrieval for a hop."""
  dropout = qa_config.dropout if is_training else 0.0
  attention_mask = modeling.create_attention_mask_from_input_mask(
      qry_input_ids, qry_input_mask)

  # question start
  with tf.variable_scope("qry/start" + suffix, reuse=tf.AUTO_REUSE):
    qry_start_seq_emb = modeling.transformer_model(
        input_tensor=qry_seq_emb,
        attention_mask=attention_mask,
        hidden_size=bert_config.hidden_size,
        num_hidden_layers=1,
        num_attention_heads=bert_config.num_attention_heads,
        intermediate_size=bert_config.intermediate_size,
        intermediate_act_fn=modeling.get_activation(bert_config.hidden_act),
        hidden_dropout_prob=dropout,
        attention_probs_dropout_prob=dropout,
        initializer_range=bert_config.initializer_range,
        do_return_all_layers=False)
    qry_start_emb = tf.squeeze(qry_start_seq_emb[:, 0:1, :], axis=1)

  # question end
  with tf.variable_scope("qry/end" + suffix, reuse=tf.AUTO_REUSE):
    qry_end_seq_emb = modeling.transformer_model(
        input_tensor=qry_seq_emb,
        attention_mask=attention_mask,
        hidden_size=bert_config.hidden_size,
        num_hidden_layers=1,
        num_attention_heads=bert_config.num_attention_heads,
        intermediate_size=bert_config.intermediate_size,
        intermediate_act_fn=modeling.get_activation(bert_config.hidden_act),
        hidden_dropout_prob=dropout,
        attention_probs_dropout_prob=dropout,
        initializer_range=bert_config.initializer_range,
        do_return_all_layers=False)
    qry_end_emb = tf.squeeze(qry_end_seq_emb[:, 0:1, :], axis=1)

  if project and qa_config.projection_dim is not None:
    with tf.variable_scope("projection"):
      qry_start_emb = contrib_layers.fully_connected(
          qry_start_emb,
          qa_config.projection_dim,
          activation_fn=tf.nn.tanh,
          reuse=tf.AUTO_REUSE,
          scope="qry_projection")
      qry_end_emb = contrib_layers.fully_connected(
          qry_end_emb,
          qa_config.projection_dim,
          activation_fn=tf.nn.tanh,
          reuse=tf.AUTO_REUSE,
          scope="qry_projection")

  return qry_start_emb, qry_end_emb
Embed query into start and end vectors for dense retrieval for a hop.
language/labs/drkit/model_fns.py