def shared_qry_encoder(qry_input_ids, qry_input_mask, is_training,
                       use_one_hot_embeddings, bert_config, qa_config):
  """Embed query into a BOW and shared dense representation."""
  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)

  # Word embeddings.
  with tf.variable_scope("bert", reuse=tf.AUTO_REUSE):
    with tf.variable_scope("embeddings", reuse=tf.AUTO_REUSE):
      # Perform embedding lookup on the word ids.
      qry_word_emb, word_emb_table = modeling.embedding_lookup(
          input_ids=qry_input_ids,
          vocab_size=bert_config.vocab_size,
          embedding_size=bert_config.hidden_size,
          initializer_range=bert_config.initializer_range,
          word_embedding_name="word_embeddings",
          use_one_hot_embeddings=use_one_hot_embeddings)

  # question shared encoder
  with tf.variable_scope("qry/encoder", reuse=tf.AUTO_REUSE):
    qry_seq_emb = modeling.transformer_model(
        input_tensor=qry_word_emb,
        attention_mask=attention_mask,
        hidden_size=bert_config.hidden_size,
        num_hidden_layers=qa_config.qry_num_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)

  return qry_seq_emb, word_emb_table
Embed query into a BOW and shared dense representation.
language/labs/drkit/model_fns.py