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=" ")
Get the prediction text for a single row in the batch.
language/question_answering/decatt_docreader/models/nq_short_pipeline_model.py