def create_objective_fn(output_layer, num_labels, obj_type, prob_vector=None):
  """Build an objective function using final layer of flex BERT model."""
  hidden_size = output_layer.shape[-1].value

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

  output_bias = tf.get_variable(
      "output_bias", [num_labels], initializer=tf.zeros_initializer())

  with tf.variable_scope("loss"):

    logits = tf.matmul(output_layer, output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias)
    probabilities = tf.nn.softmax(logits, axis=-1)
    log_probs = tf.nn.log_softmax(logits, axis=-1)

    # self-entropy of the input data-point, maximizing uncertainty
    per_example_self_entropy = -tf.reduce_sum(
        probabilities * log_probs, axis=-1)
    total_self_entropy = tf.reduce_mean(per_example_self_entropy)

    # margin of confidence which needs to be minimized for maximum uncertainty
    top_2_probs, _ = tf.math.top_k(probabilities, k=2)
    per_example_margin = top_2_probs[:, 0] - top_2_probs[:, 1]
    total_margin = tf.reduce_mean(per_example_margin)

    top_2_log_probs, _ = tf.math.top_k(log_probs, k=2)
    per_example_log_margin = top_2_log_probs[:, 0] - top_2_log_probs[:, 1]
    total_log_margin = tf.reduce_mean(per_example_log_margin)

    if prob_vector is not None:
      # Broadcasting multiplication of probability vector along batches
      per_example_cross_entropy = -tf.reduce_sum(
          prob_vector * log_probs, axis=-1)
      total_cross_entropy = tf.reduce_mean(per_example_cross_entropy)

  if "self_entropy" in obj_type:
    return total_self_entropy, per_example_self_entropy, probabilities

  elif "cross_entropy" in obj_type:
    return total_cross_entropy, per_example_cross_entropy, probabilities

  elif "confidence_margin" in obj_type:
    return total_margin, per_example_margin, probabilities

  elif "confidence_log_margin" in obj_type:
    return total_log_margin, per_example_log_margin, probabilities
Build an objective function using final layer of flex BERT model.
language/bert_extraction/steal_bert_classifier/embedding_perturbations/embedding_util.py