subword-nmt/subword_nmt/learn_joint_bpe_and_vocab.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: Rico Sennrich

"""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

#hack to get imports working if running this as a script, or within a package
if __name__ == '__main__':
    import learn_bpe
    import apply_bpe
else:
    from . import learn_bpe
    from . import apply_bpe

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('rb'), required=True, nargs = '+',
        metavar='PATH',
        help="Input texts (multiple allowed).")
    parser.add_argument(
        '--output', '-o', type=argparse.FileType('wb'), 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(
        '--byte', '-b', action="store_true",
        help="byte-level BPE.")
    parser.add_argument(
        '--separator', type=bytes, default=b'@@', metavar='STR',
        help="Separator between non-final subword units (default: '%(default)s')")
    parser.add_argument(
        '--write-vocabulary', type=argparse.FileType('wb'), 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)

    if args.byte:
        # read/write files as byte streams
        args.input = [codecs.open(f.name, 'rb') for f in args.input]
        args.vocab = [codecs.open(f.name, 'wb') for f in args.vocab]
    else:
        # read/write files as UTF-8
        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]

    args.separator = args.separator.decode('UTF-8') if not args.byte else args.separator

    # get combined vocabulary of all input texts
    full_vocab = Counter()
    for f in args.input:
        full_vocab += learn_bpe.get_vocabulary(f, num_workers=args.num_workers, is_bytes=args.byte)
        f.seek(0)

    if args.byte:
        vocab_list = [key + b' ' + str(freq).encode('UTF-8') for (key, freq) in full_vocab.items()]
    else:
        vocab_list = ['{0} {1}'.format(key, freq) for (key, freq) in full_vocab.items()]

    # learn BPE on combined vocabulary

    if args.byte:
        with open(args.output.name, 'wb') as output:
            learn_bpe.learn_bpe(vocab_list, output, args.symbols, args.min_frequency, args.verbose, is_dict=True, is_bytes=args.byte, total_symbols=args.total_symbols)

        with open(args.output.name, 'rb') as codes:
            bpe = apply_bpe.BPE(codes, separator=args.separator, is_bytes=args.byte)
    else:
        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, is_bytes=args.byte, total_symbols=args.total_symbols)

        with codecs.open(args.output.name, encoding='UTF-8') as codes:
            bpe = apply_bpe.BPE(codes, separator=args.separator, is_bytes=args.byte)

    # apply BPE to each training corpus and get vocabulary
    for train_file, vocab_file in zip(args.input, args.vocab):

        # read/write files as UTF-8
        if not args.byte:
            train_file = codecs.open(train_file.name, encoding='utf-8')
            vocab_file = codecs.open(vocab_file.name, 'w', encoding='utf-8')

        tmp = tempfile.NamedTemporaryFile(delete=False)
        tmp.close()

        if args.byte:
            tmpout = open(tmp.name, 'wb')
        else:
            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()

        if args.byte:
            tmpin = open(tmp.name, 'rb')
        else:
            tmpin = codecs.open(tmp.name, encoding='UTF-8')

        vocab = learn_bpe.get_vocabulary(tmpin, num_workers=args.num_workers, is_bytes=args.byte)
        tmpin.close()
        os.remove(tmp.name)

        for key, freq in sorted(vocab.items(), key=lambda x: x[1], reverse=True):
            if args.byte:
                vocab_file.write(key + b" " + str(freq).encode('utf-8') + b"\n")
            else:
                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
        )

    # python 2/3 compatibility
    if sys.version_info < (3, 0):
        print("Python 2 is deprecated. Use Python 3")
        sys.exit(1)

    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()

    assert(len(args.input) == len(args.vocab))

    learn_joint_bpe_and_vocab(args)