def get_text_summary(question, context, start_predictions, end_predictions):
  """Get a text summary of the question and the predicted answer."""
  question_text = tf.reduce_join(question, -1, separator=" ")

  def _get_prediction_text(args, window=5):
    """Get the prediction text for a single row in the batch."""
    current_context, start, end = args
    prediction_context_start = tf.maximum(start - window, 0)
    prediction_context_end = tf.minimum(end + 1 + window,
                                        tf.shape(current_context)[0])
    before = current_context[prediction_context_start:start]
    prediction = current_context[start:end + 1]
    after = current_context[end + 1:prediction_context_end]
    concat = tf.concat([before, ["**"], prediction, ["**"], after], 0)
    return tf.reduce_join(concat, separator=" ")

  prediction_text = tf.map_fn(
      fn=_get_prediction_text,
      elems=[context, start_predictions, end_predictions],
      dtype=tf.string,
      back_prop=False)

  return tf.summary.text("predictions",
                         tf.stack([question_text, prediction_text], -1))
Get a text summary of the question and the predicted answer.
language/question_answering/decatt_docreader/models/nq_short_pipeline_model.py