def make_binary_tag_dataset(input_prefix, output_prefix, lang,
num_workers):
logger.info("Adding tag indexes")
n_seq_tok = [0, 0]
replaced = Counter()
def merge_result(worker_result):
n_seq_tok[0] += worker_result["nseq"]
input_file = "{}{}".format(input_prefix,
("." + lang) if lang is not None else "")
offsets = Binarizer.find_offsets(input_file, num_workers)
pool = None
assert num_workers == 1
ds = indexed_dataset.make_builder(
dataset_dest_file(args, output_prefix, lang, "bin"),
impl=args.dataset_impl,
)
merge_result(
Binarizer.binarize_tag(input_file,
lambda t: ds.add_item(t),
offset=0,
end=offsets[1]))
if num_workers > 1:
pool.join()
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
temp_file_path = dataset_dest_prefix(args, prefix, lang)
ds.merge_file_(temp_file_path)
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
def make_binary_dataset(vocab, input_prefix, output_prefix, src_lang,
tgt_lang, lang, num_workers):
print("| [{}] Dictionary: {} types".format(lang, len(vocab) - 1))
n_seq_tok = [0, 0]
replaced = Counter()
def merge_result(worker_result):
replaced.update(worker_result["replaced"])
n_seq_tok[0] += worker_result["nseq"]
n_seq_tok[1] += worker_result["ntok"]
input_file = "{}.{}-{}.{}".format(input_prefix, src_lang, tgt_lang,
lang)
if args.model:
input_file += ".tok"
offsets = Binarizer.find_offsets(input_file, num_workers)
pool = None
if num_workers > 1:
pool = multiprocessing.Pool(processes=num_workers - 1)
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
pool.apply_async(
binarize,
(args, input_file, vocab, prefix, src_lang, tgt_lang, lang,
offsets[worker_id], offsets[worker_id + 1]),
callback=merge_result)
pool.close()
ds = indexed_dataset.make_builder(dataset_dest_file(
args, output_prefix, src_lang, tgt_lang, lang, "bin"),
impl=args.dataset_impl,
vocab_size=len(vocab))
merge_result(
Binarizer.binarize(input_file,
vocab,
lambda t: ds.add_item(t),
offset=0,
end=offsets[1]))
if num_workers > 1:
pool.join()
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
temp_file_path = dataset_dest_prefix(args, prefix, src_lang,
tgt_lang, lang)
ds.merge_file_(temp_file_path)
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize(
dataset_dest_file(args, output_prefix, src_lang, tgt_lang, lang,
"idx"))
print("| [{}] {}: {} sents, {} tokens, {:.3}% replaced by {}".format(
lang,
input_file,
n_seq_tok[0],
n_seq_tok[1],
100 * sum(replaced.values()) / n_seq_tok[1],
vocab.unk_word,
))
def make_binary_da_dataset(input_prefix, output_prefix, lang, num_workers,
da_mapping):
logger.info("Adding domain indexes")
n_seq_tok = [0, 0]
replaced = Counter()
def merge_result(worker_result):
n_seq_tok[0] += worker_result["nseq"]
input_file = "{}{}".format(input_prefix,
("." + lang) if lang is not None else "")
offsets = Binarizer.find_offsets(input_file, num_workers)
pool = None
# TODO: Error encounters if num_workers>1:
# No such file or directory: 'data-bin/iwslt14.tokenized.de-en/train.da1.en-de.en.idx'
if num_workers > 1:
pool = Pool(processes=num_workers - 1)
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
pool.apply_async(
binarize_da,
(
args,
input_file,
prefix,
lang,
da_mapping,
offsets[worker_id],
offsets[worker_id + 1],
),
callback=merge_result,
)
pool.close()
ds = indexed_dataset.make_builder(
dataset_dest_file(args, output_prefix, lang, "bin"),
impl=args.dataset_impl,
)
merge_result(
Binarizer.binarize_da(input_file,
lambda t: ds.add_item(t),
da_mapping,
offset=0,
end=offsets[1]))
if num_workers > 1:
pool.join()
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
temp_file_path = dataset_dest_prefix(args, prefix, lang)
ds.merge_file_(temp_file_path)
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
def make_binary_alignment_dataset(input_prefix, output_prefix,
num_workers):
nseq = [0]
def merge_result(worker_result):
nseq[0] += worker_result["nseq"]
input_file = input_prefix
offsets = Binarizer.find_offsets(input_file, num_workers)
pool = None
if num_workers > 1:
pool = Pool(processes=num_workers - 1)
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
pool.apply_async(
binarize_alignments,
(
args,
input_file,
utils.parse_alignment,
prefix,
offsets[worker_id],
offsets[worker_id + 1],
),
callback=merge_result,
)
pool.close()
ds = indexed_dataset.make_builder(dataset_dest_file(
args, output_prefix, None, "bin"),
impl=args.dataset_impl)
merge_result(
Binarizer.binarize_alignments(
input_file,
utils.parse_alignment,
lambda t: ds.add_item(t),
offset=0,
end=offsets[1],
))
if num_workers > 1:
pool.join()
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
temp_file_path = dataset_dest_prefix(args, prefix, None)
ds.merge_file_(temp_file_path)
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize(dataset_dest_file(args, output_prefix, None, "idx"))
logger.info("[alignments] {}: parsed {} alignments".format(
input_file, nseq[0]))
def make_binary_sent_doc_dataset(input_prefix, output_prefix, lang,
num_workers):
n_seq_tok = [0, 0]
replaced = Counter()
def merge_result(worker_result):
replaced.update(worker_result["replaced"])
n_seq_tok[0] += worker_result["nseq"]
n_seq_tok[1] += worker_result["ntok"]
input_file = "{}{}".format(input_prefix,
("." + lang) if lang is not None else "")
offsets = Binarizer.find_offsets(input_file, num_workers)
pool = None
if num_workers > 1:
pool = Pool(processes=num_workers - 1)
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
pool.apply_async(binarize_sent_doc,
(args, input_file, prefix, lang,
offsets[worker_id], offsets[worker_id + 1]),
callback=merge_result)
pool.close()
ds = indexed_dataset.make_builder(dataset_dest_file(
args, output_prefix, lang, "bin"),
impl=args.dataset_impl)
merge_result(
Binarizer.binarize_sent_doc(input_file,
lambda t: ds.add_item(t),
offset=0,
end=offsets[1]))
if num_workers > 1:
pool.join()
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
temp_file_path = dataset_dest_prefix(args, prefix, lang)
ds.merge_file_(temp_file_path)
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
print("| [{}] {}: {} sents, {} tokens".format(
lang,
input_file,
n_seq_tok[0],
n_seq_tok[1],
))
def binarize(args,
filename,
vocab,
output_prefix,
lang,
offset,
end,
append_eos=True,
copy_from=None):
ds = indexed_dataset.IndexedDatasetBuilder(
dataset_dest_file(args, output_prefix, lang, "bin"))
words_list = [] # todo: 目前传不出去
def consumer(ids, words):
ds.add_item(ids)
words_list.append(words)
res = Binarizer.binarize(filename,
vocab,
consumer,
append_eos=append_eos,
offset=offset,
end=end)
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
return res
def make_binary_sent_doc_dataset(input_prefix, output_prefix, lang,
num_workers, output_lang,
output_text_file):
n_seq_tok = [0, 0]
replaced = Counter()
def merge_result(worker_result):
replaced.update(worker_result["replaced"])
n_seq_tok[0] += worker_result["nseq"]
n_seq_tok[1] += worker_result["ntok"]
input_file = "{}{}".format(input_prefix,
("." + lang) if lang is not None else "")
ds = []
merge_result(
Binarizer.binarize_sent_doc(input_file, lambda t: ds.append(t)))
import json
with open(output_text_file, 'w') as f:
for line in ds:
f.write(json.dumps(line.numpy().tolist()) + '\n')
print("| [{}] {}: {} sents, {} tokens".format(
output_lang,
input_file,
n_seq_tok[0],
n_seq_tok[1],
))
def binarize(args,
filename,
vocab,
output_prefix,
lang,
offset,
end,
append_eos=True):
ds = indexed_dataset.make_builder(
dataset_dest_file(args, output_prefix, lang, "bin"),
impl=args.dataset_impl,
vocab_size=len(vocab),
)
def consumer(tensor):
ds.add_item(tensor)
res = Binarizer.binarize(filename,
vocab,
consumer,
append_eos=append_eos,
offset=offset,
end=end)
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
return res
def _make_edge_dataset(vocab, input_prefix, output_prefix, lang,
num_workers, output_text_file):
print("| [{}] Dictionary: {} types".format(lang, len(vocab) - 1))
n_seq_tok = [0, 0]
replaced = Counter()
def merge_result(worker_result):
replaced.update(worker_result["replaced"])
n_seq_tok[0] += worker_result["nseq"]
n_seq_tok[1] += worker_result["ntok"]
input_file = "{}{}".format(input_prefix,
("." + lang) if lang is not None else "")
ds = []
merge_result(
Binarizer.binarize_graph(input_file, vocab,
lambda t: ds.append(t)))
import json
with open(output_text_file, 'w') as f:
for line in ds:
f.write(json.dumps(line.numpy().tolist()) + '\n')
print("| [{}] {}: {} sents, {} tokens, {:.3}% replaced by {}".format(
lang,
input_file,
n_seq_tok[0],
n_seq_tok[1],
100 * sum(replaced.values()) / n_seq_tok[1],
vocab.unk_word,
))
def binarize(args,
filename,
vocab,
output_prefix,
lang,
offset,
end,
append_eos=True):
ds = indexed_dataset.make_builder(dataset_dest_file(
args, output_prefix, lang, "bin"),
impl=args.dataset_impl,
vocab_size=len(vocab))
def consumer(tensor):
ds.add_item(tensor)
tk = tokenize_smiles if args.file_format == 'smiles' else tokenize_line
res = Binarizer.binarize(filename,
vocab,
consumer,
tokenize=tk,
append_eos=append_eos,
offset=offset,
end=end)
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
return res
def make_binary_alignment_dataset(input_prefix, output_prefix, src, tgt,
num_workers):
nseq = [0]
def merge_result(worker_result):
nseq[0] += worker_result['nseq']
parse_alignment = lambda s: torch.IntTensor(
[int(t) for t in s.split()])
input_file = input_prefix
offsets = Binarizer.find_offsets(input_file, num_workers)
pool = None
if num_workers > 1:
pool = multiprocessing.Pool(processes=num_workers - 1)
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
pool.apply_async(
binarize_alignments,
(args, input_file, parse_alignment, prefix, src, tgt,
offsets[worker_id], offsets[worker_id + 1]),
callback=merge_result)
pool.close()
ds = indexed_dataset.make_builder(dataset_dest_file(
args, output_prefix, src, tgt, None, "bin"),
impl=args.dataset_impl)
merge_result(
Binarizer.binarize_alignments(input_file,
parse_alignment,
lambda t: ds.add_item(t),
offset=0,
end=offsets[1]))
if num_workers > 1:
pool.join()
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
temp_file_path = dataset_dest_prefix(args, prefix, src, tgt)
ds.merge_file_(temp_file_path)
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize(
dataset_dest_file(args, output_prefix, src, tgt, None, "idx"))
print("| [alignments] {}: parsed {} alignments".format(
input_file, nseq[0]))
def binarize(args, filename, dict, output_prefix, lang, offset, end):
ds = indexed_dataset.IndexedDatasetBuilder(dataset_dest_file(args, output_prefix, lang, 'bin'))
def consumer(tensor):
ds.add_item(tensor)
res = Binarizer.binarize(filename, dict, consumer, offset=offset, end=end)
ds.finalize(dataset_dest_file(args, output_prefix, lang, 'idx'))
return res
def make_binary_dataset(input_prefix, output_prefix, lang, num_workers):
dict = dictionary.Dictionary.load(dict_path(lang))
print('| [{}] Dictionary: {} types'.format(lang, len(dict) - 1))
n_seq_tok = [0, 0]
replaced = Counter()
def merge_result(worker_result):
replaced.update(worker_result['replaced'])
n_seq_tok[0] += worker_result['nseq']
n_seq_tok[1] += worker_result['ntok']
input_file = '{}{}'.format(input_prefix, ('.' + lang) if lang is not None else '')
offsets = Binarizer.find_offsets(input_file, num_workers)
pool = None
if num_workers > 1:
pool = Pool(processes=num_workers-1)
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
pool.apply_async(binarize, (args, input_file, dict, prefix, lang,
offsets[worker_id],
offsets[worker_id + 1]), callback=merge_result)
pool.close()
ds = indexed_dataset.IndexedDatasetBuilder(dataset_dest_file(args, output_prefix, lang, 'bin'))
merge_result(Binarizer.binarize(input_file, dict, lambda t: ds.add_item(t),
offset=0, end=offsets[1]))
if num_workers > 1:
pool.join()
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
temp_file_path = dataset_dest_prefix(args, prefix, lang)
ds.merge_file_(temp_file_path)
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize(dataset_dest_file(args, output_prefix, lang, 'idx'))
print('| [{}] {}: {} sents, {} tokens, {:.3}% replaced by {}'.format(
lang, input_file, n_seq_tok[0], n_seq_tok[1],
100 * sum(replaced.values()) / n_seq_tok[1], dict.unk_word))
def binarize_alignments(args, filename, parse_alignment, output_prefix, offset, end):
ds = indexed_dataset.make_builder(dataset_dest_file(args, output_prefix, None, "bin"),
impl=args.dataset_impl, vocab_size=None)
def consumer(tensor):
ds.add_item(tensor)
res = Binarizer.binarize_alignments(filename, parse_alignment, consumer, offset=offset,
end=end)
ds.finalize(dataset_dest_file(args, output_prefix, None, "idx"))
return res
def binarize(args, filename, vocab, output_prefix, lang, offset, end,
tgt_dict):
append_eos = True
ds = indexed_dataset.IndexedDatasetBuilder(
dataset_dest_file(args, output_prefix, lang, "bin"))
def consumer(tensor):
ds.add_item(tensor)
res = Binarizer.binarize(filename,
vocab,
consumer,
append_eos=append_eos,
reverse_order=args.reverse_order,
offset=offset,
end=end,
tgt_dict=tgt_dict)
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
return res
def binarize(args,
filename,
dict,
output_prefix,
lang,
offset,
end,
append_eos=True):
ds = indexed_dataset.IndexedDatasetBuilder(
dataset_dest_file(args, output_prefix, lang, "bin"))
def consumer(tensor):
ds.add_item(tensor)
res = Binarizer.binarize(filename,
dict,
consumer,
offset=offset,
end=end,
append_eos=append_eos)
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
return res
def binarize(args,
filename,
vocab,
output_prefix,
lang,
offset,
end,
append_eos=True):
##dataset_impl=mmap, ds -> MMapIndexedDatasetBuilder
##dataset_impl=lazy, ds -> IndexedDatasetBuilder
ds = indexed_dataset.make_builder(dataset_dest_file(
args, output_prefix, lang, "bin"),
impl=args.dataset_impl,
vocab_size=len(vocab))
def consumer(tensor):
##输入的tensor，就是直接把文本串通过dictionary转为id串
##dataset_impl=mmap, MMapIndexedDatasetBuilder.add_item:把输入tensor直接写入文件
##dataset_impl=lazy, IndexedDatasetBuilder.add_item:把输入tensor写入文件，并更新sizes, data_offsets, dim_offsets
ds.add_item(tensor)
## 读入文件filename在offset和end之间的内容，并把每个文本串利用dictionary转为id串，利用consumer函数写入到ds中
res = Binarizer.binarize(filename,
vocab,
consumer,
append_eos=append_eos,
offset=offset,
end=end)
##把写入到ds中的数据存储到对应路径的临时文件, output_prefix包含了worker_id，以区分不同的worker的临时文件
##mmap, ds.finalize:调用MMapIndexedDataset.write写入三个tensor:
## 训练样例数量，每个样例tensor的size，每个样例的位置pointer
##lazy, ds.finalize:IndexedDatasetBuilder.finalize直接写入dim_offsets, data_offsets, sizes
# data_offsets: 存放每个tensor在二进制文件中的结尾位置（前一个tensor的结尾处就是这个tensor的开始位置）
# sizes: 存放每个tensor的shape的各个dim值
# dim_offsets: 存放每个tensor的shape在self.size中的结尾位置（前面tensor shape的结尾是这个tensor shape的开始）
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
return res
def get_offsets(input_file, num_workers):
return Binarizer.find_offsets(input_file, num_workers)
def make_binary_dataset(vocab, input_prefix, output_prefix, lang,
num_workers):
logger.info("[{}] Dictionary: {} types".format(lang, len(vocab)))
output_prefix += '.bert' if isinstance(vocab, BertTokenizer) else ''
input_prefix += '.bert' if isinstance(vocab, BertTokenizer) else ''
n_seq_tok = [0, 0]
replaced = Counter()
def merge_result(worker_result):
replaced.update(worker_result["replaced"])
n_seq_tok[0] += worker_result["nseq"]
n_seq_tok[1] += worker_result["ntok"]
input_file = "{}{}".format(input_prefix,
("." + lang) if lang is not None else "")
offsets = Binarizer.find_offsets(input_file, num_workers)
pool = None
if num_workers > 1:
pool = Pool(processes=num_workers - 1)
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
pool.apply_async(
binarize,
(
args,
input_file,
vocab,
prefix,
lang,
offsets[worker_id],
offsets[worker_id + 1],
),
callback=merge_result,
)
pool.close()
ds = indexed_dataset.make_builder(
dataset_dest_file(args, output_prefix, lang, "bin"),
impl=args.dataset_impl,
vocab_size=len(vocab),
)
merge_result(
Binarizer.binarize(input_file,
vocab,
lambda t: ds.add_item(t),
offset=0,
end=offsets[1]))
if num_workers > 1:
pool.join()
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
temp_file_path = dataset_dest_prefix(args, prefix, lang)
ds.merge_file_(temp_file_path)
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
logger.info(
"[{}] {}: {} sents, {} tokens, {:.3}% replaced by {}".format(
lang,
input_file,
n_seq_tok[0],
n_seq_tok[1],
100 * sum(replaced.values()) / n_seq_tok[1],
vocab.unk_word,
))
def make_binary_dataset(vocab, input_prefix, output_prefix, lang, num_workers, avoid_tokenize=False):
if vocab is not None:
print("| [{}] Dictionary: {} types".format(lang, len(vocab) - 1))
else:
print('| Using None Dictionary and only string split is performed.')
n_seq_tok = [0, 0]
replaced = Counter()
def merge_result(worker_result):
replaced.update(worker_result["replaced"])
n_seq_tok[0] += worker_result["nseq"]
n_seq_tok[1] += worker_result["ntok"]
input_file = "{}{}".format(
input_prefix, ("." + lang) if lang is not None else ""
)
offsets = Binarizer.find_offsets(input_file, num_workers)
pool = None
if num_workers > 1:
pool = Pool(processes=num_workers - 1)
for worker_id in range(1, num_workers):
# TODO: worker > 1 is not working for map dataset
if args.input_mapping is True:
raise NotImplementedError("Worker > 1 is not implemented for map dataset yet.")
prefix = "{}{}".format(output_prefix, worker_id)
pool.apply_async(
binarize,
(
args,
input_file,
vocab,
prefix,
lang,
offsets[worker_id],
offsets[worker_id + 1],
avoid_tokenize,
),
callback=merge_result,
)
pool.close()
ds = indexed_dataset.make_builder(
dataset_dest_file(args, output_prefix, lang, "bin"),
impl=args.dataset_impl,
vocab_size=len(vocab) if vocab is not None else -1,
)
merge_result(
Binarizer.binarize(
input_file, vocab, lambda t: ds.add_item(t), offset=0, end=offsets[1], avoid_tokenize=avoid_tokenize,
)
)
if num_workers > 1:
pool.join()
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
temp_file_path = dataset_dest_prefix(args, prefix, lang)
ds.merge_file_(temp_file_path)
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
if vocab is not None:
unk = vocab.unk_word if hasattr(vocab, 'unk_word') else vocab.unk_token
else:
unk = ""
logger.info(
"[{}] {}: {} sents, {} tokens, {:.3}% replaced by {}".format(
lang,
input_file,
n_seq_tok[0],
n_seq_tok[1],
100 * sum(replaced.values()) / n_seq_tok[1],
unk,
)
)
def main(args):
print(args)
os.makedirs(args.destdir, exist_ok=True)
target = not args.only_source
def build_dictionary(filenames):
d = dictionary.Dictionary()
for filename in filenames:
dictionary.Dictionary.add_file_to_dictionary(filename, d, tokenize_line, args.workers)
return d
def train_path(lang):
return '{}{}'.format(args.trainpref, ('.' + lang) if lang else '')
def file_name(prefix, lang):
fname = prefix
if lang is not None:
fname += f'.{lang}'
return fname
def dest_path(prefix, lang):
return os.path.join(args.destdir, file_name(prefix, lang))
def dict_path(lang):
return dest_path('dict', lang) + '.txt'
if args.joined_dictionary:
assert not args.srcdict, 'cannot combine --srcdict and --joined-dictionary'
assert not args.tgtdict, 'cannot combine --tgtdict and --joined-dictionary'
src_dict = build_dictionary(set([
train_path(lang)
for lang in [args.source_lang, args.target_lang]
]))
tgt_dict = src_dict
else:
if args.srcdict:
src_dict = dictionary.Dictionary.load(args.srcdict)
else:
assert args.trainpref, "--trainpref must be set if --srcdict is not specified"
src_dict = build_dictionary([train_path(args.source_lang)])
if target:
if args.tgtdict:
tgt_dict = dictionary.Dictionary.load(args.tgtdict)
else:
assert args.trainpref, "--trainpref must be set if --tgtdict is not specified"
tgt_dict = build_dictionary([train_path(args.target_lang)])
src_dict.finalize(
threshold=args.thresholdsrc,
nwords=args.nwordssrc,
padding_factor=args.padding_factor,
)
src_dict.save(dict_path(args.source_lang))
if target:
if not args.joined_dictionary:
tgt_dict.finalize(
threshold=args.thresholdtgt,
nwords=args.nwordstgt,
padding_factor=args.padding_factor,
)
tgt_dict.save(dict_path(args.target_lang))
def make_binary_dataset(input_prefix, output_prefix, lang, num_workers):
dict = dictionary.Dictionary.load(dict_path(lang))
print('| [{}] Dictionary: {} types'.format(lang, len(dict) - 1))
n_seq_tok = [0, 0]
replaced = Counter()
def merge_result(worker_result):
replaced.update(worker_result['replaced'])
n_seq_tok[0] += worker_result['nseq']
n_seq_tok[1] += worker_result['ntok']
input_file = '{}{}'.format(input_prefix, ('.' + lang) if lang is not None else '')
offsets = Binarizer.find_offsets(input_file, num_workers)
pool = None
if num_workers > 1:
pool = Pool(processes=num_workers-1)
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
pool.apply_async(binarize, (args, input_file, dict, prefix, lang,
offsets[worker_id],
offsets[worker_id + 1]), callback=merge_result)
pool.close()
ds = indexed_dataset.IndexedDatasetBuilder(dataset_dest_file(args, output_prefix, lang, 'bin'))
merge_result(Binarizer.binarize(input_file, dict, lambda t: ds.add_item(t),
offset=0, end=offsets[1]))
if num_workers > 1:
pool.join()
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
temp_file_path = dataset_dest_prefix(args, prefix, lang)
ds.merge_file_(temp_file_path)
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize(dataset_dest_file(args, output_prefix, lang, 'idx'))
print('| [{}] {}: {} sents, {} tokens, {:.3}% replaced by {}'.format(
lang, input_file, n_seq_tok[0], n_seq_tok[1],
100 * sum(replaced.values()) / n_seq_tok[1], dict.unk_word))
def make_dataset(input_prefix, output_prefix, lang, num_workers=1):
if args.output_format == 'binary':
make_binary_dataset(input_prefix, output_prefix, lang, num_workers)
elif args.output_format == 'raw':
# Copy original text file to destination folder
output_text_file = dest_path(
output_prefix + '.{}-{}'.format(args.source_lang, args.target_lang),
lang,
)
shutil.copyfile(file_name(input_prefix, lang), output_text_file)
def make_all(lang):
if args.trainpref:
make_dataset(args.trainpref, 'train', lang, num_workers=args.workers)
if args.validpref:
for k, validpref in enumerate(args.validpref.split(',')):
outprefix = 'valid{}'.format(k) if k > 0 else 'valid'
make_dataset(validpref, outprefix, lang)
if args.testpref:
for k, testpref in enumerate(args.testpref.split(',')):
outprefix = 'test{}'.format(k) if k > 0 else 'test'
make_dataset(testpref, outprefix, lang)
make_all(args.source_lang)
if target:
make_all(args.target_lang)
print('| Wrote preprocessed data to {}'.format(args.destdir))
if args.alignfile:
assert args.trainpref, "--trainpref must be set if --alignfile is specified"
src_file_name = train_path(args.source_lang)
tgt_file_name = train_path(args.target_lang)
src_dict = dictionary.Dictionary.load(dict_path(args.source_lang))
tgt_dict = dictionary.Dictionary.load(dict_path(args.target_lang))
freq_map = {}
with open(args.alignfile, 'r') as align_file:
with open(src_file_name, 'r') as src_file:
with open(tgt_file_name, 'r') as tgt_file:
for a, s, t in zip_longest(align_file, src_file, tgt_file):
si = Binarizer.tokenize(s, src_dict, add_if_not_exist=False)
ti = Binarizer.tokenize(t, tgt_dict, add_if_not_exist=False)
ai = list(map(lambda x: tuple(x.split('-')), a.split()))
for sai, tai in ai:
srcidx = si[int(sai)]
tgtidx = ti[int(tai)]
if srcidx != src_dict.unk() and tgtidx != tgt_dict.unk():
assert srcidx != src_dict.pad()
assert srcidx != src_dict.eos()
assert tgtidx != tgt_dict.pad()
assert tgtidx != tgt_dict.eos()
if srcidx not in freq_map:
freq_map[srcidx] = {}
if tgtidx not in freq_map[srcidx]:
freq_map[srcidx][tgtidx] = 1
else:
freq_map[srcidx][tgtidx] += 1
align_dict = {}
for srcidx in freq_map.keys():
align_dict[srcidx] = max(freq_map[srcidx], key=freq_map[srcidx].get)
with open(os.path.join(args.destdir, 'alignment.{}-{}.txt'.format(
args.source_lang, args.target_lang)), 'w') as f:
for k, v in align_dict.items():
print('{} {}'.format(src_dict[k], tgt_dict[v]), file=f)
def make_binary_dataset(vocab, input_prefix, output_prefix, lang, num_workers, copy_src_words=None):
print("| [{}] Dictionary: {} types".format(lang, len(vocab) - 1))
n_seq_tok = [0, 0]
replaced = Counter()
copyied = Counter()
def merge_result(worker_result):
replaced.update(worker_result["replaced"])
copyied.update(worker_result["copied"])
n_seq_tok[0] += worker_result["nseq"]
n_seq_tok[1] += worker_result["ntok"]
input_file = "{}{}".format(
input_prefix, ("." + lang) if lang is not None else ""
)
offsets = Binarizer.find_offsets(input_file, num_workers)
pool = None
if num_workers > 1: # todo: not support copy
pool = Pool(processes=num_workers - 1)
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
pool.apply_async(
binarize,
(
args,
input_file,
vocab,
prefix,
lang,
offsets[worker_id],
offsets[worker_id + 1]
),
callback=merge_result
)
pool.close()
ds = indexed_dataset.IndexedDatasetBuilder(
dataset_dest_file(args, output_prefix, lang, "bin")
)
words_list = []
def binarize_consumer(ids, words):
ds.add_item(ids)
words_list.append(words)
merge_result(
Binarizer.binarize(
input_file, vocab, binarize_consumer,
offset=0, end=offsets[1], copy_ext_dict=args.copy_ext_dict, copy_src_words=copy_src_words
)
)
if num_workers > 1:
pool.join()
for worker_id in range(1, num_workers):
prefix = "{}{}".format(output_prefix, worker_id)
temp_file_path = dataset_dest_prefix(args, prefix, lang)
ds.merge_file_(temp_file_path)
os.remove(indexed_dataset.data_file_path(temp_file_path))
os.remove(indexed_dataset.index_file_path(temp_file_path))
ds.finalize(dataset_dest_file(args, output_prefix, lang, "idx"))
print(
"| [{}] {}: {} sents, {} tokens, {:.3}% replaced by {}, {:.3}% <unk> copied from src".format(
lang,
input_file,
n_seq_tok[0],
n_seq_tok[1],
100 * sum(replaced.values()) / n_seq_tok[1],
vocab.unk_word,
100 * sum(copyied.values()) / n_seq_tok[1]
)
)
return words_list
