def model_fn(features, labels, mode, params):
  """A model function satisfying the tf.estimator API."""
  del labels
  assert mode == tf_estimator.ModeKeys.PREDICT, "Mode should be PREDICT."
  rc_model = hub.Module(params["rc_model"])
  outputs = rc_model(
      inputs=dict(
          input_ids=tf.cast(features["input_ids"], tf.int32),
          input_mask=tf.cast(features["input_mask"], tf.int32),
          segment_ids=tf.cast(features["segment_ids"], tf.int32)),
      signature="extractive_qa",
      as_dict=True)
  start, end, _ = reward_utils.max_scoring_span(
      start_scores=outputs["start_logits"],
      end_scores=outputs["end_logits"],
      max_length=params["max_answer_length"],
      no_answer_bias=params["no_answer_bias"])
  is_consistent = tf.logical_and(
      tf.logical_and(tf.greater(start, 0), tf.greater(end, 0)),
      tf.logical_and(
          tf.equal(start, tf.cast(features["start_positions"] - 1, tf.int32)),
          tf.equal(end, tf.cast(features["end_positions"] - 1, tf.int32))))
  return tf_estimator.tpu.TPUEstimatorSpec(
      mode=mode,
      predictions=dict(
          unique_ids=features["unique_ids"],
          input_ids=features["input_ids"],
          start=start,
          end=end,
          is_consistent=is_consistent))
A model function satisfying the tf.estimator API.
language/capwap/synthetic/filter_round_trip.py