def model_fn(input_tensor, bert_input_mask, token_type_ids, bert_config,
             num_labels, obj_type, prob_vector):
  """Create a flex input model and decide the objective function."""
  flex_bert_model = BertFlexEmbeddingModel(
      config=bert_config,
      is_training=False,
      input_tensor=input_tensor,
      input_mask=bert_input_mask,
      token_type_ids=token_type_ids)

  output_layer = flex_bert_model.get_pooled_output()

  objective, per_eg_obj, probs = create_objective_fn(
      output_layer=output_layer,
      num_labels=num_labels,
      obj_type=obj_type,
      prob_vector=prob_vector)
  return objective, per_eg_obj, probs
Create a flex input model and decide the objective function.
language/bert_extraction/steal_bert_classifier/embedding_perturbations/embedding_util.py