def build_save_in_shards(src_corpus, tgt_corpus, fields,
corpus_type, opt, ref_corpus=None):
corpus_size = os.path.getsize(src_corpus)
if corpus_size > 10 * (1024 ** 2) and opt.max_shard_size == 0:
logger.info("Warning. The corpus %s is larger than 10M bytes, "
"you can set '-max_shard_size' to process it by "
"small shards to use less memory." % src_corpus)
if opt.max_shard_size != 0:
logger.info(' * divide corpus into shards and build dataset '
'separately (shard_size = %d bytes).'
% opt.max_shard_size)
ret_list = []
src_iter = inputters.ShardedTextCorpusIterator(
src_corpus, opt.src_seq_length_trunc,
"src", opt.max_shard_size)
tgt_iter = inputters.ShardedTextCorpusIterator(
tgt_corpus, opt.tgt_seq_length_trunc,
"tgt", opt.max_shard_size,
assoc_iter=src_iter)
if ref_corpus:
ref_iter = inputters.ShardedTextCorpusIterator(
ref_corpus, opt.tgt_seq_length_trunc,
"ref", opt.max_shard_size,
assoc_iter=src_iter)
else:
ref_iter = None
index = 0
while not src_iter.hit_end():
index += 1
dataset = inputters.TextDataset(
fields, src_iter, tgt_iter,
src_iter.num_feats, tgt_iter.num_feats,
src_seq_length=opt.src_seq_length,
tgt_seq_length=opt.tgt_seq_length,
dynamic_dict=opt.dynamic_dict, ref_examples_iter=ref_iter,
num_ref_feats=ref_iter.num_feats if ref_iter else None)
# We save fields in vocab.pt separately, so make it empty.
dataset.fields = []
pt_file = "{:s}.{:s}.{:d}.pt".format(
opt.save_data, corpus_type, index)
logger.info(" * saving %s data shard to %s."
% (corpus_type, pt_file))
torch.save(dataset, pt_file)
ret_list.append(pt_file)
return ret_list
def build_save_in_shards(src_corpus, tgt_corpus, fields, corpus_type, opt):
"""
Divide the big corpus into shards, and build dataset separately.
This is currently only for data_type=='text'.
The reason we do this is to avoid taking up too much memory due
to sucking in a huge corpus file.
To tackle this, we only read in part of the corpus file of size
`max_shard_size`(actually it is multiples of 64 bytes that equals
or is slightly larger than this size), and process it into dataset,
then write it to disk along the way. By doing this, we only focus on
part of the corpus at any moment, thus effectively reducing memory use.
According to test, this method can reduce memory footprint by ~50%.
Note! As we process along the shards, previous shards might still
stay in memory, but since we are done with them, and no more
reference to them, if there is memory tight situation, the OS could
easily reclaim these memory.
If `max_shard_size` is 0 or is larger than the corpus size, it is
effectively preprocessed into one dataset, i.e. no sharding.
NOTE! `max_shard_size` is measuring the input corpus size, not the
output pt file size. So a shard pt file consists of examples of size
2 * `max_shard_size`(source + target).
"""
corpus_size = os.path.getsize(src_corpus)
if corpus_size > 10 * (1024**2) and opt.max_shard_size == 0:
logger.info("Warning. The corpus %s is larger than 10M bytes, "
"you can set '-max_shard_size' to process it by "
"small shards to use less memory." % src_corpus)
if opt.max_shard_size != 0:
logger.info(' * divide corpus into shards and build dataset '
'separately (shard_size = %d bytes).' % opt.max_shard_size)
ret_list = []
src_iter = inputters.ShardedTextCorpusIterator(src_corpus,
opt.src_seq_length_trunc,
"src", opt.max_shard_size)
tgt_iter = inputters.ShardedTextCorpusIterator(tgt_corpus,
opt.tgt_seq_length_trunc,
"tgt",
opt.max_shard_size,
assoc_iter=src_iter)
index = 0
while not src_iter.hit_end():
index += 1
dataset = inputters.TextDataset(fields,
src_iter,
tgt_iter,
src_iter.num_feats,
tgt_iter.num_feats,
src_seq_length=opt.src_seq_length,
tgt_seq_length=opt.tgt_seq_length,
dynamic_dict=opt.dynamic_dict)
# We save fields in vocab.pt separately, so make it empty.
dataset.fields = []
pt_file = "{:s}.{:s}.{:d}.pt".format(opt.save_data, corpus_type, index)
logger.info(" * saving %s data shard to %s." % (corpus_type, pt_file))
torch.save(dataset, pt_file)
ret_list.append(pt_file)
return ret_list
