def create_de_model(bert_config, is_training, input_ids_1, input_mask_1,
                    segment_ids_1, input_ids_2, input_masks_2, segment_ids_2,
                    num_candidates, labels, use_one_hot_embeddings):
  """Creates a ranking model using cosine and dual encoder representations."""

  sequence_length_query = FLAGS.max_seq_length_query
  sequence_length_passage = FLAGS.max_seq_length - FLAGS.max_seq_length_query

  input_ids_1 = tf.reshape(input_ids_1, [-1, sequence_length_query])
  segment_ids_1 = tf.reshape(segment_ids_1, [-1, sequence_length_query])
  input_masks_1 = tf.reshape(input_mask_1, [-1, sequence_length_query])
  batch_size = tf.shape(input_masks_1)[0]

  input_ids_2 = tf.reshape(input_ids_2, [-1, sequence_length_passage])
  segment_ids_2 = tf.reshape(segment_ids_2, [-1, sequence_length_passage])
  input_masks_2 = tf.reshape(input_masks_2, [-1, sequence_length_passage])

  # [batch_size, num_candidates]
  labels = tf.dtypes.cast(labels, tf.float32)

  # [batch_size, num_vec_query, hidden_size], [batch_size, num_vec_query]
  output_layer_1, mask_1 = encode_block(bert_config, input_ids_1, input_masks_1,
                                        segment_ids_1, use_one_hot_embeddings,
                                        FLAGS.num_vec_query, is_training)

  output_layer_2, mask_2 = encode_block(bert_config, input_ids_2, input_masks_2,
                                        segment_ids_2, use_one_hot_embeddings,
                                        FLAGS.num_vec_passage, is_training)

  label_mask = tf.expand_dims(tf.eye(batch_size), axis=2)
  label_mask = tf.tile(label_mask, [1, 1, num_candidates])
  label_mask = tf.reshape(label_mask, [batch_size, -1])
  label_mask = tf.cast(label_mask, tf.float32)

  labels = tf.tile(labels, [1, batch_size])
  labels = tf.multiply(labels, label_mask)
  output_layer_2_logits = tf.reshape(
      output_layer_2, [batch_size, num_candidates, FLAGS.num_vec_passage, -1])
  mask_2_logits = tf.reshape(
      mask_2, [batch_size, num_candidates, FLAGS.num_vec_passage])
  mask_logits = tf.einsum("BQ,BCP->BCQP", tf.cast(mask_1, tf.float32),
                          tf.cast(mask_2_logits, tf.float32))

  logits = tf.einsum("BQH,BCPH->BCQP", output_layer_1, output_layer_2_logits)
  logits = tf.multiply(logits, mask_logits)
  logits = tf.reduce_max(logits, axis=-1)
  logits = tf.reduce_sum(logits, axis=-1)

  if FLAGS.use_tpu and is_training:
    num_shards = tpu_utils.num_tpu_shards()
    output_layer_2 = tpu_utils.cross_shard_concat(output_layer_2)
    mask_2 = tpu_utils.cross_shard_concat(tf.cast(mask_2, tf.float32))
    mask_2 = tf.cast(mask_2, tf.bool)
    labels = tpu_utils.cross_shard_pad(labels)
    tf.logging.info("Global batch size: %s", tensor_utils.shape(labels, 0))
    tf.logging.info("Num shards: %s", num_shards)
    tf.logging.info("Number of candidates in batch: %s",
                    tensor_utils.shape(output_layer_2, 0))
    labels = tf.reshape(labels, [num_shards, batch_size, -1])
    labels = tf.transpose(labels, perm=[1, 0, 2])
    labels = tf.reshape(labels, [batch_size, -1])

  with tf.variable_scope("loss"):
    if is_training:
      output_layer_1 = tf.nn.dropout(output_layer_1, keep_prob=FLAGS.dropout)
      output_layer_2 = tf.nn.dropout(output_layer_2, keep_prob=FLAGS.dropout)
    cosine_similarity = tf.einsum("AQH,BPH->ABQP", output_layer_1,
                                  output_layer_2)
    mask = tf.cast(
        tf.logical_and(
            tf.expand_dims(tf.expand_dims(mask_1, 2), 1),
            tf.expand_dims(tf.expand_dims(mask_2, 1), 0)), tf.float32)
    cosine_similarity = tf.multiply(cosine_similarity, mask)
    cosine_similarity = tf.reduce_max(cosine_similarity, axis=-1)
    cosine_similarity = tf.reduce_sum(cosine_similarity, axis=-1)
    per_example_loss = tf.losses.softmax_cross_entropy(labels,
                                                       cosine_similarity)

    return (per_example_loss, logits)
Creates a ranking model using cosine and dual encoder representations.
language/multivec/models/ranking_model_experiment_inbatch.py