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
def get_pretrain_dataset ( data , batch_size , num_ctxes , shuffle , use_avg_len , num_buckets , num_parts = 1 , part_idx = 0 , prefetch = True ) : 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 ) : lengths = dataset . get_field ( 'valid_lengths' ) batchify_fn = Tuple ( Pad ( ) , Pad ( ) , Pad ( ) , Pad ( ) , Stack ( ) , Pad ( ) , Stack ( ) ) if use_avg_len : sampler = nlp . data . FixedBucketSampler ( lengths = lengths , 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.