def build_dataset(task):
    """Helper function to select and build dataset."""

    if task == 'LAMBADA':
        return _build_lambada_dataset()
    if task == 'WIKITEXT103':
        return _build_wikitext103_dataset()

    raise NotImplementedError('dataset for {} task is not '
                              'implemented.'.format(task))
Helper function to select and build dataset.
megatron_lm/tasks/zeroshot_gpt2/datasets.py