def dummy_input(config: ml_collections.ConfigDict) -> Dict[Text, Any]:
    """Produces model-specific dummy input batch. See BaseTask for details."""

    if config.get('max_length_with_entity_tokens') is not None:
      max_length = config.max_length_with_entity_tokens
    else:
      max_length = config.model_config.encoder_config.max_length

    bsz = config.per_device_batch_size
    text_shape = (bsz, max_length)
    mention_shape = (config.max_mentions)
    mention_target_shape = (2 * bsz)
    int_type = jnp.int32

    position_ids = np.arange(max_length)
    position_ids = np.tile(position_ids, (bsz, 1))

    dummy_input = {
        'text_ids':
            jnp.ones(text_shape, int_type),
        'text_mask':
            jnp.ones(text_shape, int_type),
        'position_ids':
            jnp.asarray(position_ids, int_type),
        'segment_ids':
            jnp.zeros(text_shape, int_type),
        'classifier_target':
            jnp.ones((bsz,), int_type),
        'mention_start_positions':
            jnp.zeros(mention_shape, int_type),
        'mention_end_positions':
            jnp.zeros(mention_shape, int_type),
        'mention_mask':
            jnp.ones(mention_shape, int_type),
        'mention_batch_positions':
            jnp.ones(mention_shape, int_type),
        'mention_target_indices':
            jnp.arange(mention_target_shape, dtype=int_type),
        'mention_target_weights':
            jnp.ones(mention_target_shape, dtype=int_type),
        'mention_object_indices':
            jnp.arange(bsz, dtype=int_type),
        'mention_subject_indices':
            jnp.arange(bsz, dtype=int_type),
        'mention_target_batch_positions':
            jnp.arange(mention_target_shape, dtype=int_type),
        'mention_target_start_positions':
            jnp.zeros(mention_target_shape, int_type),
        'mention_target_end_positions':
            jnp.zeros(mention_target_shape, int_type),
        'mention_target_ids':
            jnp.zeros(mention_target_shape, int_type),
    }

    return dummy_input
Produces model-specific dummy input batch. See BaseTask for details.
language/mentionmemory/tasks/relation_classifier_task.py