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."
    span, _ = rc_span(
        question_ids=features["question_inputs"].token_ids,
        question_mask=features["question_inputs"].mask,
        context_ids=features["context_inputs"].token_ids,
        context_mask=features["context_inputs"].mask,
        rc_model=hub.Module(params["rc_model"]),
        vocab=vocab,
        no_answer_bias=params.get("no_answer_bias", -1e4))
    return tf_estimator.tpu.TPUEstimatorSpec(
        mode=mode,
        predictions=dict(
            question_id=features["question_id"],
            image_id=features["image_id"],
            beam_id=features["beam_id"],
            span=span))
A model function satisfying the tf.estimator API.
language/capwap/utils/reward_utils.py