def preprocess_mapper(features, params, lookup_table, vocab, mode):
  """Model-specific preprocessing of features from the dataset."""
  # Set input type.
  features["input_type"] = tf.constant(datasets.DatasetTypes.REFERENCE)

  if mode != tf_estimator.ModeKeys.PREDICT:
    # Select random caption.
    captions = tf.io.parse_tensor(features["captions"], tf.string)
    num_captions = tensor_utils.shape(captions, 0)
    rid = tf.random.uniform([], maxval=num_captions, dtype=tf.int32)

    caption = text_utils.build_text_inputs(
        text=captions[rid],
        length=params["caption_length"],
        lookup_table=lookup_table,
        segment_id=0,
        start_token=vocab.CLS,
        end_token=vocab.SEP)
    assert isinstance(caption, text_utils.TextInputs)

    features["token_inputs"] = text_utils.TextInputs(
        token_ids=caption.token_ids[:-1],
        mask=caption.mask[:-1],
        segment_ids=caption.segment_ids[:-1],
        positions=caption.positions[:-1])

    features["token_outputs"] = text_utils.TextOutputs(
        token_ids=caption.token_ids[1:], mask=caption.mask[1:])

    if params.get("conditional_decoding"):
      random_span = text_utils.get_random_span(
          text=captions[rid],
          p=params["span_sample_p"],
          max_span_len=params["span_length"])

      features["condition_inputs"] = text_utils.build_text_inputs(
          text=random_span,
          length=params["condition_length"],
          lookup_table=lookup_table,
          segment_id=1,
          start_token=vocab.ANS)

  features["object_features"] = image_utils.parse_object_features(
      features["object_features"], features["object_positions"], params)

  # Remove extra inputs.
  features = {f: features[f] for f in features if f in KEYS}

  # Add dummy inputs for standardization for multi-tasking.
  footprint = datasets.footprint(params)
  assert footprint
  for k, v in footprint.items():
    if k not in features:
      features[k] = v

  return features
Model-specific preprocessing of features from the dataset.
language/capwap/datasets/captions_dataset.py