def get_name_to_features(config: ml_collections.ConfigDict) -> Dict[str, Any]:
    """Return feature dict for decoding purposes. See BaseTask."""

    encoder_config = config.model_config.encoder_config
    max_length = encoder_config.max_length
    name_to_features = {
        'text_ids': tf.io.FixedLenFeature(max_length, tf.int64),
        'text_mask': tf.io.FixedLenFeature(max_length, tf.int64),
        'dense_span_starts': tf.io.FixedLenFeature(max_length, tf.int64),
        'dense_span_ends': tf.io.FixedLenFeature(max_length, tf.int64),
        'dense_mention_mask': tf.io.FixedLenFeature(max_length, tf.int64),
        'dense_mention_ids': tf.io.FixedLenFeature(max_length, tf.int64),
    }

    return name_to_features
Return feature dict for decoding purposes. See BaseTask.
language/mentionmemory/tasks/mention_encoder_task.py