def build_sample(row_id, context_ids, context_types, context_pad_mask):
    """Convert to numpy and return a sample consumed by the batch producer."""

    context_ids = np.array(context_ids, dtype=np.int64)
    context_types = np.array(context_types, dtype=np.int64)
    context_mask = make_attention_mask(context_ids, context_ids)

    sample = ({
        'row_id': row_id,
        'context': context_ids,
        'context_mask': context_mask,
        'context_types': context_types,
        'context_pad_mask': context_pad_mask
    })
    return sample
Convert to numpy and return a sample consumed by the batch producer.
training/DeepSpeed-Domino/megatron/data/orqa_wiki_dataset.py