wehub-resource-sync
327 lines
12 KiB
Python
327 lines
12 KiB
Python
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import argparse
|
|
import os
|
|
import shutil
|
|
from itertools import zip_longest
|
|
from pprint import pprint
|
|
|
|
from paddlenlp.data import Vocab
|
|
from paddlenlp.utils.log import logger
|
|
|
|
|
|
def get_preprocessing_parser():
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument("-s", "--src_lang", default=None, type=str, help="Source language. ")
|
|
parser.add_argument("-t", "--trg_lang", default=None, type=str, help="Target language. ")
|
|
parser.add_argument(
|
|
"--train_pref", default=None, type=str, help="The prefix for train file and also used to save dict. "
|
|
)
|
|
parser.add_argument(
|
|
"--dev_pref",
|
|
default=None,
|
|
type=str,
|
|
help="The prefixes for dev file and use comma to separate. "
|
|
"(words missing from train set are replaced with <unk>)",
|
|
)
|
|
parser.add_argument(
|
|
"--test_pref",
|
|
default=None,
|
|
type=str,
|
|
help="The prefixes for test file and use comma to separate. "
|
|
"(words missing from train set are replaced with <unk>)",
|
|
)
|
|
parser.add_argument(
|
|
"--dest_dir",
|
|
default="./data/",
|
|
type=str,
|
|
help="The destination dir to save processed train, dev and test file. ",
|
|
)
|
|
parser.add_argument(
|
|
"--threshold_trg", default=0, type=int, help="Map words appearing less than threshold times to unknown. "
|
|
)
|
|
parser.add_argument(
|
|
"--threshold_src", default=0, type=int, help="Map words appearing less than threshold times to unknown. "
|
|
)
|
|
parser.add_argument("--src_vocab", default=None, type=str, help="Reuse given source dictionary. ")
|
|
parser.add_argument("--trg_vocab", default=None, type=str, help="Reuse given target dictionary. ")
|
|
parser.add_argument("--nwords_trg", default=None, type=int, help="The number of target words to retain. ")
|
|
parser.add_argument("--nwords_src", default=None, type=int, help="The number of source words to retain. ")
|
|
parser.add_argument("--align_file", default=None, help="An alignment file (optional). ")
|
|
parser.add_argument("--joined_dictionary", action="store_true", help="Generate joined dictionary. ")
|
|
parser.add_argument("--only_source", action="store_true", help="Only process the source language. ")
|
|
parser.add_argument(
|
|
"--dict_only", action="store_true", help="Only builds a dictionary and then exits if it's set."
|
|
)
|
|
parser.add_argument("--bos_token", default="<s>", type=str, help="bos_token. ")
|
|
parser.add_argument("--eos_token", default="</s>", type=str, help="eos_token. ")
|
|
parser.add_argument(
|
|
"--pad_token",
|
|
default=None,
|
|
type=str,
|
|
help="The token used for padding. If it's None, the bos_token will be used. Defaults to None. ",
|
|
)
|
|
parser.add_argument("--unk_token", default="<unk>", type=str, help="Unk token. ")
|
|
parser.add_argument("--apply_bpe", action="store_true", help="Whether to apply bpe to the files. ")
|
|
parser.add_argument(
|
|
"--bpe_code", default=None, type=str, help="The code used for bpe. Must be provided when --apply_bpe is set. "
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def _train_path(lang, train_pref):
|
|
return "{}{}".format(train_pref, ("." + lang) if lang else "")
|
|
|
|
|
|
def _dev_path(lang, dev_pref):
|
|
return "{}{}".format(dev_pref, ("." + lang) if lang else "")
|
|
|
|
|
|
def _test_path(lang, test_pref):
|
|
return "{}{}".format(test_pref, ("." + lang) if lang else "")
|
|
|
|
|
|
def _file_name(prefix, lang):
|
|
fname = prefix
|
|
if lang is not None:
|
|
fname += ".{lang}".format(lang=lang)
|
|
return fname
|
|
|
|
|
|
def _dest_path(prefix, lang, dest_dir):
|
|
return os.path.join(dest_dir, _file_name(prefix, lang))
|
|
|
|
|
|
def _dict_path(lang, dest_dir):
|
|
return _dest_path("dict", lang, dest_dir) + ".txt"
|
|
|
|
|
|
def _build_dictionary(filenames, args, src=False, trg=False):
|
|
assert src ^ trg, "src and trg cannot be both True or both False. "
|
|
|
|
if not isinstance(filenames, (list, tuple)):
|
|
filenames = [filenames]
|
|
|
|
tokens = []
|
|
for file in filenames:
|
|
with open(file, "r") as f:
|
|
lines = f.readlines()
|
|
for line in lines:
|
|
tokens.append(line.strip().split())
|
|
|
|
return Vocab.build_vocab(
|
|
tokens,
|
|
max_size=args.nwords_src if src else args.nwords_trg,
|
|
min_freq=args.threshold_src if src else args.threshold_trg,
|
|
unk_token=args.unk_token,
|
|
pad_token=args.pad_token,
|
|
bos_token=args.bos_token,
|
|
eos_token=args.eos_token,
|
|
)
|
|
|
|
|
|
def _make_dataset(vocab, input_prefix, output_prefix, lang, args):
|
|
# Copy original text file to destination folder
|
|
output_text_file = _dest_path(
|
|
output_prefix + ".{}-{}".format(args.src_lang, args.trg_lang),
|
|
lang,
|
|
args.dest_dir,
|
|
)
|
|
|
|
shutil.copyfile(_file_name(input_prefix, lang), output_text_file)
|
|
|
|
|
|
def _make_all(lang, vocab, args):
|
|
if args.train_pref:
|
|
_make_dataset(vocab, args.train_pref, "train", lang, args=args)
|
|
|
|
if args.dev_pref:
|
|
for k, dev_pref in enumerate(args.dev_pref.split(",")):
|
|
out_prefix = "dev{}".format(k) if k > 0 else "dev"
|
|
_make_dataset(vocab, dev_pref, out_prefix, lang, args=args)
|
|
|
|
if args.test_pref:
|
|
for k, test_pref in enumerate(args.test_pref.split(",")):
|
|
out_prefix = "test{}".format(k) if k > 0 else "test"
|
|
_make_dataset(vocab, test_pref, out_prefix, lang, args=args)
|
|
|
|
|
|
def _align_files(args, src_vocab, trg_vocab):
|
|
assert args.train_pref, "--train_pref must be set if --align_file is specified"
|
|
src_file_name = _train_path(args.src_lang, args.train_pref)
|
|
trg_file_name = _train_path(args.trg_lang, args.train_pref)
|
|
freq_map = {}
|
|
|
|
with open(args.align_file, "r", encoding="utf-8") as align_file:
|
|
with open(src_file_name, "r", encoding="utf-8") as src_file:
|
|
with open(trg_file_name, "r", encoding="utf-8") as trg_file:
|
|
for a, s, t in zip_longest(align_file, src_file, trg_file):
|
|
si = src_vocab.to_indices(s)
|
|
ti = trg_vocab.to_indices(t)
|
|
ai = list(map(lambda x: tuple(x.split("\t")), a.split()))
|
|
for sai, tai in ai:
|
|
src_idx = si[int(sai)]
|
|
trg_idx = ti[int(tai)]
|
|
if src_idx != src_vocab.get_unk_token_id() and trg_idx != trg_vocab.get_unk_token_id():
|
|
assert src_idx != src_vocab.get_pad_token_id()
|
|
assert src_idx != src_vocab.get_eos_token_id()
|
|
assert trg_idx != trg_vocab.get_pad_token_id()
|
|
assert trg_idx != trg_vocab.get_eos_token_id()
|
|
if src_idx not in freq_map:
|
|
freq_map[src_idx] = {}
|
|
if trg_idx not in freq_map[src_idx]:
|
|
freq_map[src_idx][trg_idx] = 1
|
|
else:
|
|
freq_map[src_idx][trg_idx] += 1
|
|
|
|
align_dict = {}
|
|
for src_idx in freq_map.keys():
|
|
align_dict[src_idx] = max(freq_map[src_idx], key=freq_map[src_idx].get)
|
|
|
|
with open(
|
|
os.path.join(
|
|
args.dest_dir,
|
|
"alignment.{}-{}.txt".format(args.src_lang, args.trg_lang),
|
|
),
|
|
"w",
|
|
encoding="utf-8",
|
|
) as f:
|
|
for k, v in align_dict.items():
|
|
print("{} {}".format(src_vocab[k], trg_vocab[v]), file=f)
|
|
|
|
|
|
def main(args):
|
|
os.makedirs(args.dest_dir, exist_ok=True)
|
|
pprint(args)
|
|
|
|
if args.apply_bpe:
|
|
import fastBPE
|
|
|
|
bpe = fastBPE.fastBPE(args.bpe_code)
|
|
filenames = [_train_path(lang, args.train_pref) for lang in [args.src_lang, args.trg_lang]]
|
|
for k, dev_pref in enumerate(args.dev_pref.split(",")):
|
|
filenames.extend([_dev_path(lang, args.dev_pref) for lang in [args.src_lang, args.trg_lang]])
|
|
for k, test_pref in enumerate(args.test_pref.split(",")):
|
|
filenames.extend([_test_path(lang, args.test_pref) for lang in [args.src_lang, args.trg_lang]])
|
|
|
|
for file in filenames:
|
|
sequences = []
|
|
with open(file, "r") as f:
|
|
lines = f.readlines()
|
|
for seq in lines:
|
|
sequences.append(seq.strip())
|
|
|
|
bpe_sequences = bpe.apply(sequences)
|
|
os.makedirs(os.path.join(args.train_pref, "tmp_bpe"), exist_ok=True)
|
|
shutil.copyfile(file, os.path.join(args.train_pref, "tmp_bpe", os.path.split(file)[-1]))
|
|
|
|
with open(file, "w") as f:
|
|
for bpe_seq in bpe_sequences:
|
|
f.write(bpe_seq + "\n")
|
|
|
|
# build dictionaries
|
|
target = not args.only_source
|
|
|
|
if not args.src_vocab and os.path.exists(_dict_path(args.src_lang, args.dest_dir)):
|
|
raise FileExistsError(_dict_path(args.src_lang, args.dest_dir))
|
|
|
|
if target and not args.trg_vocab and os.path.exists(_dict_path(args.trg_lang, args.dest_dir)):
|
|
raise FileExistsError(_dict_path(args.trg_lang, args.dest_dir))
|
|
|
|
if args.joined_dictionary:
|
|
assert (
|
|
not args.src_vocab or not args.trg_vocab
|
|
), "Cannot use both --src_vocab and --trg_vocab with --joined_dictionary"
|
|
|
|
if args.src_vocab:
|
|
src_vocab = Vocab.load_vocabulary(
|
|
filepath=args.src_vocab,
|
|
unk_token=args.unk_token,
|
|
bos_token=args.bos_token,
|
|
eos_token=args.eos_token,
|
|
pad_token=args.pad_token,
|
|
)
|
|
elif args.trg_vocab:
|
|
src_vocab = Vocab.load_vocabulary(
|
|
filepath=args.trg_vocab,
|
|
unk_token=args.unk_token,
|
|
bos_token=args.bos_token,
|
|
eos_token=args.eos_token,
|
|
pad_token=args.pad_token,
|
|
)
|
|
else:
|
|
assert args.train_pref, "--train_pref must be set if --src_vocab is not specified. "
|
|
src_vocab = _build_dictionary(
|
|
[_train_path(lang, args.train_pref) for lang in [args.src_lang, args.trg_lang]], args=args, src=True
|
|
)
|
|
|
|
trg_vocab = src_vocab
|
|
else:
|
|
if args.src_vocab:
|
|
src_vocab = Vocab.load_vocabulary(
|
|
filepath=args.src_vocab,
|
|
unk_token=args.unk_token,
|
|
bos_token=args.bos_token,
|
|
eos_token=args.eos_token,
|
|
pad_token=args.pad_token,
|
|
)
|
|
else:
|
|
assert args.train_pref, "--train_pref must be set if --src_vocab is not specified"
|
|
src_vocab = _build_dictionary([_train_path(args.src_lang, args.train_pref)], args=args, src=True)
|
|
|
|
if target:
|
|
if args.trg_vocab:
|
|
trg_vocab = Vocab.load_vocabulary(
|
|
filepath=args.trg_vocab,
|
|
unk_token=args.unk_token,
|
|
bos_token=args.bos_token,
|
|
eos_token=args.eos_token,
|
|
pad_token=args.pad_token,
|
|
)
|
|
else:
|
|
assert args.train_pref, "--train_pref must be set if --trg_vocab is not specified"
|
|
trg_vocab = _build_dictionary([_train_path(args.trg_lang, args.train_pref)], args=args, trg=True)
|
|
else:
|
|
trg_vocab = None
|
|
|
|
# save dictionaries
|
|
src_vocab.save_vocabulary(_dict_path(args.src_lang, args.dest_dir))
|
|
if target and trg_vocab is not None:
|
|
trg_vocab.save_vocabulary(_dict_path(args.trg_lang, args.dest_dir))
|
|
|
|
if args.dict_only:
|
|
return
|
|
|
|
_make_all(args.src_lang, src_vocab, args)
|
|
if target:
|
|
_make_all(args.trg_lang, trg_vocab, args)
|
|
|
|
logger.info("Wrote preprocessed data to {}".format(args.dest_dir))
|
|
|
|
if args.align_file:
|
|
_align_files(args, src_vocab=src_vocab, trg_vocab=trg_vocab)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = get_preprocessing_parser()
|
|
main(args)
|