def get_auxiliary_output(cls, loss_helpers: Dict[str, Any]) -> Dict[str, Any]:
    """Extract features from `loss_helpers` to be used as `auxiliary_output`."""
    auxiliary_output = {}

    # As part of auxiliary output for the downstream tasks we save retrieved
    # memory information if it exists in the `loss_helper`.
    # This extra information would allow us to manually analyze models
    # performance on the task.

    def add_memory_auxiliary_output(prefix=''):
      keys = ('top_entity_ids', 'top_memory_ids', 'memory_attention_weights')
      for key in keys:
        full_key = prefix + key
        if full_key in loss_helpers:
          auxiliary_output[full_key] = loss_helpers[full_key]

    add_memory_auxiliary_output()
    add_memory_auxiliary_output('second_')

    return auxiliary_output
Extract features from `loss_helpers` to be used as `auxiliary_output`.
language/mentionmemory/tasks/downstream_encoder_task.py