def get_pretrain_dataset(data, batch_size, num_ctxes, shuffle, use_avg_len,
                         num_buckets, num_parts=1, part_idx=0, prefetch=True):
    """create dataset for pretraining."""
    num_files = len(glob.glob(os.path.expanduser(data)))
    logging.debug('%d files found.', num_files)
    assert num_files >= num_parts, \
        'Number of training files must be greater than the number of partitions'
    split_sampler = nlp.data.SplitSampler(num_files, num_parts=num_parts, part_index=part_idx)
    stream = nlp.data.SimpleDatasetStream(nlp.data.NumpyDataset, data, split_sampler)
    if prefetch:
        stream = nlp.data.PrefetchingStream(stream)

    def get_dataloader(dataset):
        """create data loader based on the dataset chunk"""
        lengths = dataset.get_field('valid_lengths')
        # A batch includes: input_id, masked_id, masked_position, masked_weight,
        #                   next_sentence_label, segment_id, valid_length
        batchify_fn = Tuple(Pad(), Pad(), Pad(), Pad(), Stack(), Pad(), Stack())
        if use_avg_len:
            # sharded data loader
            sampler = nlp.data.FixedBucketSampler(lengths=lengths,
                                                  # batch_size per shard
                                                  batch_size=batch_size,
                                                  num_buckets=num_buckets,
                                                  shuffle=shuffle,
                                                  use_average_length=True,
                                                  num_shards=num_ctxes)
            dataloader = nlp.data.ShardedDataLoader(dataset,
                                                    batch_sampler=sampler,
                                                    batchify_fn=batchify_fn,
                                                    num_workers=num_ctxes)
        else:
            sampler = nlp.data.FixedBucketSampler(lengths,
                                                  batch_size=batch_size * num_ctxes,
                                                  num_buckets=num_buckets,
                                                  ratio=0,
                                                  shuffle=shuffle)
            dataloader = DataLoader(dataset=dataset,
                                    batch_sampler=sampler,
                                    batchify_fn=batchify_fn,
                                    num_workers=1)
        logging.debug('Sampler created for a new dataset:\n%s', sampler.stats())
        return dataloader

    stream = stream.transform(get_dataloader)
    return stream
create dataset for pretraining.