def get_pretrain_dataset(data, batch_size, num_ctxes, shuffle, use_avg_len,<EOL>num_buckets, num_parts=<NUM_LIT:1>, part_idx=<NUM_LIT:0>, prefetch=True):
num_files = len(glob.glob(os.path.expanduser(data)))<EOL>logging.debug('<STR_LIT>', num_files)<EOL>assert num_files >= num_parts,'<STR_LIT>'<EOL>split_sampler = nlp.data.SplitSampler(num_files, num_parts=num_parts, part_index=part_idx)<EOL>stream = nlp.data.SimpleDatasetStream(nlp.data.NumpyDataset, data, split_sampler)<EOL>if prefetch:<EOL><INDENT>stream = nlp.data.PrefetchingStream(stream)<EOL><DEDENT>def get_dataloader(dataset):<EOL><INDENT>"""<STR_LIT>"""<EOL>lengths = dataset.get_field('<STR_LIT>')<EOL>batchify_fn = Tuple(Pad(), Pad(), Pad(), Pad(), Stack(), Pad(), Stack())<EOL>if use_avg_len:<EOL><INDENT>sampler = nlp.data.FixedBucketSampler(lengths=lengths,<EOL>batch_size=batch_size,<EOL>num_buckets=num_buckets,<EOL>shuffle=shuffle,<EOL>use_average_length=True,<EOL>num_shards=num_ctxes)<EOL>dataloader = nlp.data.ShardedDataLoader(dataset,<EOL>batch_sampler=sampler,<EOL>batchify_fn=batchify_fn,<EOL>num_workers=num_ctxes)<EOL><DEDENT>else:<EOL><INDENT>sampler = nlp.data.FixedBucketSampler(lengths,<EOL>batch_size=batch_size * num_ctxes,<EOL>num_buckets=num_buckets,<EOL>ratio=<NUM_LIT:0>,<EOL>shuffle=shuffle)<EOL>dataloader = DataLoader(dataset=dataset,<EOL>batch_sampler=sampler,<EOL>batchify_fn=batchify_fn,<EOL>num_workers=<NUM_LIT:1>)<EOL><DEDENT>logging.debug('<STR_LIT>', sampler.stats())<EOL>return dataloader<EOL><DEDENT>stream = stream.transform(get_dataloader)<EOL>return stream<EOL>
create dataset for pretraining.