def create_ca_model(bert_config, is_training, input_ids, input_mask,
                    segment_ids, num_candidates, sequence_length, labels,
                    use_one_hot_embeddings):
  """Creates a ranking model using cross attention representations."""

  input_ids = tf.reshape(input_ids, [-1, sequence_length])
  segment_ids = tf.reshape(segment_ids, [-1, sequence_length])
  input_mask = tf.reshape(input_mask, [-1, sequence_length])
  labels = tf.dtypes.cast(labels, tf.float32)
  model = modeling.BertModel(
      config=bert_config,
      is_training=is_training,
      input_ids=input_ids,
      input_mask=input_mask,
      token_type_ids=segment_ids,
      use_one_hot_embeddings=use_one_hot_embeddings)
  output_layer = model.get_pooled_output()

  hidden_size = output_layer.shape[-1]

  output_weights = tf.get_variable(
      "output_weights", [1, hidden_size],
      initializer=tf.truncated_normal_initializer(mean=0, stddev=0.02))

  with tf.variable_scope("loss"):
    if is_training:
      output_layer = tf.nn.dropout(output_layer, keep_prob=FLAGS.dropout)

    logits = tf.matmul(output_layer, output_weights, transpose_b=True)
    logits = tf.reshape(logits, [-1, num_candidates])
    log_probs = tf.nn.log_softmax(logits, axis=-1)
    per_example_loss = -tf.reduce_sum(labels * log_probs, axis=-1)
    loss = tf.reduce_mean(per_example_loss)
    return (loss, logits)
Creates a ranking model using cross attention representations.
language/multivec/models/ranking_model_experiment_inbatch.py