-
Notifications
You must be signed in to change notification settings - Fork 464
/
learn_joint_bpe_and_vocab.py
executable file
·191 lines (154 loc) · 7.6 KB
/
learn_joint_bpe_and_vocab.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
#!/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)