def get_name_to_features(
      config: ml_collections.ConfigDict) -> Dict[Text, Any]:
    """Return feature dict for decoding purposes. See BaseTask for details."""
    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),
        'target':
            tf.io.FixedLenFeature(config.max_num_labels_per_sample, tf.int64),
        'target_mask':
            tf.io.FixedLenFeature(config.max_num_labels_per_sample, tf.int64),
        'mention_start_positions':
            tf.io.FixedLenFeature(config.max_mentions_per_sample, tf.int64),
        'mention_end_positions':
            tf.io.FixedLenFeature(config.max_mentions_per_sample, tf.int64),
        'mention_mask':
            tf.io.FixedLenFeature(config.max_mentions_per_sample, tf.int64),
        'mention_target_indices':
            tf.io.FixedLenFeature(1, tf.int64),
    }

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