|
|
|
|
|
|
|
|
|
"""Use byte pair encoding (BPE) to learn a variable-length encoding of the vocabulary in a text. |
|
This script learns BPE jointly on a concatenation of a list of texts (typically the source and target side of a parallel corpus, |
|
applies the learned operation to each and (optionally) returns the resulting vocabulary of each text. |
|
The vocabulary can be used in apply_bpe.py to avoid producing symbols that are rare or OOV in a training text. |
|
|
|
Reference: |
|
Rico Sennrich, Barry Haddow and Alexandra Birch (2016). Neural Machine Translation of Rare Words with Subword Units. |
|
Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (ACL 2016). Berlin, Germany. |
|
""" |
|
|
|
from __future__ import unicode_literals |
|
|
|
import sys |
|
import os |
|
import inspect |
|
import codecs |
|
import argparse |
|
import tempfile |
|
import warnings |
|
from collections import Counter |
|
from multiprocessing import cpu_count |
|
|
|
|
|
if __name__ == '__main__': |
|
import learn_bpe |
|
import apply_bpe |
|
else: |
|
from . import learn_bpe |
|
from . import apply_bpe |
|
|
|
|
|
from io import open |
|
argparse.open = open |
|
|
|
def create_parser(subparsers=None): |
|
|
|
if subparsers: |
|
parser = subparsers.add_parser('learn-joint-bpe-and-vocab', |
|
formatter_class=argparse.RawDescriptionHelpFormatter, |
|
description="learn BPE-based word segmentation") |
|
else: |
|
parser = argparse.ArgumentParser( |
|
formatter_class=argparse.RawDescriptionHelpFormatter, |
|
description="learn BPE-based word segmentation") |
|
|
|
parser.add_argument( |
|
'--input', '-i', type=argparse.FileType('r'), required=True, nargs = '+', |
|
metavar='PATH', |
|
help="Input texts (multiple allowed).") |
|
parser.add_argument( |
|
'--output', '-o', type=argparse.FileType('w'), required=True, |
|
metavar='PATH', |
|
help="Output file for BPE codes.") |
|
parser.add_argument( |
|
'--symbols', '-s', type=int, default=10000, |
|
help="Create this many new symbols (each representing a character n-gram) (default: %(default)s)") |
|
parser.add_argument( |
|
'--separator', type=str, default='@@', metavar='STR', |
|
help="Separator between non-final subword units (default: '%(default)s')") |
|
parser.add_argument( |
|
'--write-vocabulary', type=argparse.FileType('w'), required=True, nargs = '+', default=None, |
|
metavar='PATH', dest='vocab', |
|
help='Write to these vocabulary files after applying BPE. One per input text. Used for filtering in apply_bpe.py') |
|
parser.add_argument( |
|
'--min-frequency', type=int, default=2, metavar='FREQ', |
|
help='Stop if no symbol pair has frequency >= FREQ (default: %(default)s)') |
|
parser.add_argument( |
|
'--total-symbols', '-t', action="store_true", |
|
help="subtract number of characters from the symbols to be generated (so that '--symbols' becomes an estimate for the total number of symbols needed to encode text).") |
|
parser.add_argument( |
|
'--num-workers', type=int, default=1, |
|
help="Number of processors to process texts, only supported in Python3. If -1, set `multiprocessing.cpu_count()`. (default: %(default)s)") |
|
parser.add_argument( |
|
'--verbose', '-v', action="store_true", |
|
help="verbose mode.") |
|
|
|
return parser |
|
|
|
def learn_joint_bpe_and_vocab(args): |
|
|
|
if args.vocab and len(args.input) != len(args.vocab): |
|
sys.stderr.write('Error: number of input files and vocabulary files must match\n') |
|
sys.exit(1) |
|
|
|
|
|
args.input = [codecs.open(f.name, encoding='UTF-8') for f in args.input] |
|
args.vocab = [codecs.open(f.name, 'w', encoding='UTF-8') for f in args.vocab] |
|
|
|
|
|
full_vocab = Counter() |
|
for f in args.input: |
|
full_vocab += learn_bpe.get_vocabulary(f, num_workers=args.num_workers) |
|
f.seek(0) |
|
|
|
vocab_list = ['{0} {1}'.format(key, freq) for (key, freq) in full_vocab.items()] |
|
|
|
|
|
with codecs.open(args.output.name, 'w', encoding='UTF-8') as output: |
|
learn_bpe.learn_bpe(vocab_list, output, args.symbols, args.min_frequency, args.verbose, is_dict=True, total_symbols=args.total_symbols) |
|
|
|
with codecs.open(args.output.name, encoding='UTF-8') as codes: |
|
bpe = apply_bpe.BPE(codes, separator=args.separator) |
|
|
|
|
|
for train_file, vocab_file in zip(args.input, args.vocab): |
|
|
|
tmp = tempfile.NamedTemporaryFile(delete=False) |
|
tmp.close() |
|
|
|
tmpout = codecs.open(tmp.name, 'w', encoding='UTF-8') |
|
|
|
train_file.seek(0) |
|
bpe.process_lines(train_file.name, tmpout, num_workers=args.num_workers) |
|
|
|
tmpout.close() |
|
tmpin = codecs.open(tmp.name, encoding='UTF-8') |
|
|
|
vocab = learn_bpe.get_vocabulary(tmpin, num_workers=args.num_workers) |
|
tmpin.close() |
|
os.remove(tmp.name) |
|
|
|
for key, freq in sorted(vocab.items(), key=lambda x: x[1], reverse=True): |
|
vocab_file.write("{0} {1}\n".format(key, freq)) |
|
train_file.close() |
|
vocab_file.close() |
|
|
|
|
|
if __name__ == '__main__': |
|
|
|
currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) |
|
newdir = os.path.join(currentdir, 'subword_nmt') |
|
if os.path.isdir(newdir): |
|
warnings.warn( |
|
"this script's location has moved to {0}. This symbolic link will be removed in a future version. Please point to the new location, or install the package and use the command 'subword-nmt'".format(newdir), |
|
DeprecationWarning |
|
) |
|
|
|
|
|
if sys.version_info < (3, 0): |
|
sys.stderr = codecs.getwriter('UTF-8')(sys.stderr) |
|
sys.stdout = codecs.getwriter('UTF-8')(sys.stdout) |
|
sys.stdin = codecs.getreader('UTF-8')(sys.stdin) |
|
else: |
|
sys.stderr = codecs.getwriter('UTF-8')(sys.stderr.buffer) |
|
sys.stdout = codecs.getwriter('UTF-8')(sys.stdout.buffer) |
|
sys.stdin = codecs.getreader('UTF-8')(sys.stdin.buffer) |
|
|
|
parser = create_parser() |
|
args = parser.parse_args() |
|
|
|
if args.num_workers <= 0: |
|
args.num_workers = cpu_count() |
|
|
|
if sys.version_info < (3, 0): |
|
args.separator = args.separator.decode('UTF-8') |
|
if args.num_workers > 1: |
|
args.num_workers = 1 |
|
warnings.warn("Parallel mode is only supported in Python3. Using 1 processor instead.") |
|
|
|
assert(len(args.input) == len(args.vocab)) |
|
|
|
learn_joint_bpe_and_vocab(args) |
|
|