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
192
193
194
195
196
197
198
199
200
201
202
203
|
"""
Entry point for training and evaluating a neural tokenizer.
This tokenizer treats tokenization and sentence segmentation as a tagging problem, and uses a combination of
recurrent and convolutional architectures.
For details please refer to paper: https://nlp.stanford.edu/pubs/qi2018universal.pdf.
"""
import argparse
from copy import copy
import logging
import random
import numpy as np
import torch
from stanza.models.common import utils
from stanza.models.tokenization.trainer import Trainer
from stanza.models.tokenization.data import DataLoader
from stanza.models.tokenization.utils import load_mwt_dict, eval_model, output_predictions
from stanza.models import _training_logging
logger = logging.getLogger('stanza')
def parse_args(args=None):
"""
If args == None, the system args are used.
"""
parser = argparse.ArgumentParser()
parser.add_argument('--txt_file', type=str, help="Input plaintext file")
parser.add_argument('--label_file', type=str, default=None, help="Character-level label file")
parser.add_argument('--json_file', type=str, default=None, help="JSON file with pre-chunked units")
parser.add_argument('--mwt_json_file', type=str, default=None, help="JSON file for MWT expansions")
parser.add_argument('--conll_file', type=str, default=None, help="CoNLL file for output")
parser.add_argument('--dev_txt_file', type=str, help="(Train only) Input plaintext file for the dev set")
parser.add_argument('--dev_label_file', type=str, default=None, help="(Train only) Character-level label file for the dev set")
parser.add_argument('--dev_json_file', type=str, default=None, help="(Train only) JSON file with pre-chunked units for the dev set")
parser.add_argument('--dev_conll_gold', type=str, default=None, help="(Train only) CoNLL-U file for the dev set for early stopping")
parser.add_argument('--lang', type=str, help="Language")
parser.add_argument('--shorthand', type=str, help="UD treebank shorthand")
parser.add_argument('--mode', default='train', choices=['train', 'predict'])
parser.add_argument('--skip_newline', action='store_true', help="Whether to skip newline characters in input. Particularly useful for languages like Chinese.")
parser.add_argument('--emb_dim', type=int, default=32, help="Dimension of unit embeddings")
parser.add_argument('--hidden_dim', type=int, default=64, help="Dimension of hidden units")
parser.add_argument('--conv_filters', type=str, default="1,9", help="Configuration of conv filters. ,, separates layers and , separates filter sizes in the same layer.")
parser.add_argument('--no-residual', dest='residual', action='store_false', help="Add linear residual connections")
parser.add_argument('--no-hierarchical', dest='hierarchical', action='store_false', help="\"Hierarchical\" RNN tokenizer")
parser.add_argument('--hier_invtemp', type=float, default=0.5, help="Inverse temperature used in propagating tokenization predictions between RNN layers")
parser.add_argument('--input_dropout', action='store_true', help="Dropout input embeddings as well")
parser.add_argument('--conv_res', type=str, default=None, help="Convolutional residual layers for the RNN")
parser.add_argument('--rnn_layers', type=int, default=1, help="Layers of RNN in the tokenizer")
parser.add_argument('--max_grad_norm', type=float, default=1.0, help="Maximum gradient norm to clip to")
parser.add_argument('--anneal', type=float, default=.999, help="Anneal the learning rate by this amount when dev performance deteriorate")
parser.add_argument('--anneal_after', type=int, default=2000, help="Anneal the learning rate no earlier than this step")
parser.add_argument('--lr0', type=float, default=2e-3, help="Initial learning rate")
parser.add_argument('--dropout', type=float, default=0.33, help="Dropout probability")
parser.add_argument('--unit_dropout', type=float, default=0.33, help="Unit dropout probability")
parser.add_argument('--tok_noise', type=float, default=0.02, help="Probability to induce noise to the input of the higher RNN")
parser.add_argument('--weight_decay', type=float, default=0.0, help="Weight decay")
parser.add_argument('--max_seqlen', type=int, default=100, help="Maximum sequence length to consider at a time")
parser.add_argument('--batch_size', type=int, default=32, help="Batch size to use")
parser.add_argument('--epochs', type=int, default=10, help="Total epochs to train the model for")
parser.add_argument('--steps', type=int, default=50000, help="Steps to train the model for, if unspecified use epochs")
parser.add_argument('--report_steps', type=int, default=20, help="Update step interval to report loss")
parser.add_argument('--shuffle_steps', type=int, default=100, help="Step interval to shuffle each paragraph in the generator")
parser.add_argument('--eval_steps', type=int, default=200, help="Step interval to evaluate the model on the dev set for early stopping")
parser.add_argument('--max_steps_before_stop', type=int, default=5000, help='Early terminates after this many steps if the dev scores are not improving')
parser.add_argument('--save_name', type=str, default=None, help="File name to save the model")
parser.add_argument('--load_name', type=str, default=None, help="File name to load a saved model")
parser.add_argument('--save_dir', type=str, default='saved_models/tokenize', help="Directory to save models in")
parser.add_argument('--cuda', type=bool, default=torch.cuda.is_available())
parser.add_argument('--cpu', action='store_true', help='Ignore CUDA and run on CPU.')
parser.add_argument('--seed', type=int, default=1234)
parser.add_argument('--use_mwt', dest='use_mwt', default=None, action='store_true', help='Whether or not to include mwt output layers. If set to None, this will be determined by examining the training data for MWTs')
parser.add_argument('--no_use_mwt', dest='use_mwt', action='store_false', help='Whether or not to include mwt output layers')
args = parser.parse_args(args=args)
return args
def main(args=None):
args = parse_args(args=args)
if args.cpu:
args.cuda = False
utils.set_random_seed(args.seed, args.cuda)
args = vars(args)
logger.info("Running tokenizer in {} mode".format(args['mode']))
args['feat_funcs'] = ['space_before', 'capitalized', 'all_caps', 'numeric']
args['feat_dim'] = len(args['feat_funcs'])
args['save_name'] = "{}/{}".format(args['save_dir'], args['save_name']) if args['save_name'] is not None \
else '{}/{}_tokenizer.pt'.format(args['save_dir'], args['shorthand'])
utils.ensure_dir(args['save_dir'])
if args['mode'] == 'train':
train(args)
else:
evaluate(args)
def train(args):
mwt_dict = load_mwt_dict(args['mwt_json_file'])
train_input_files = {
'json': args['json_file'],
'txt': args['txt_file'],
'label': args['label_file']
}
train_batches = DataLoader(args, input_files=train_input_files)
vocab = train_batches.vocab
args['vocab_size'] = len(vocab)
dev_input_files = {
'json': args['dev_json_file'],
'txt': args['dev_txt_file'],
'label': args['dev_label_file']
}
dev_batches = DataLoader(args, input_files=dev_input_files, vocab=vocab, evaluation=True)
if args['use_mwt'] is None:
args['use_mwt'] = train_batches.has_mwt()
logger.info("Found {}mwts in the training data. Setting use_mwt to {}".format(("" if args['use_mwt'] else "no "), args['use_mwt']))
trainer = Trainer(args=args, vocab=vocab, use_cuda=args['cuda'])
if args['load_name'] is not None:
load_name = "{}/{}".format(args['save_dir'], args['load_name'])
trainer.load(load_name)
trainer.change_lr(args['lr0'])
N = len(train_batches)
steps = args['steps'] if args['steps'] is not None else int(N * args['epochs'] / args['batch_size'] + .5)
lr = args['lr0']
prev_dev_score = -1
best_dev_score = -1
best_dev_step = -1
for step in range(1, steps+1):
batch = train_batches.next(unit_dropout=args['unit_dropout'])
loss = trainer.update(batch)
if step % args['report_steps'] == 0:
logger.info("Step {:6d}/{:6d} Loss: {:.3f}".format(step, steps, loss))
if args['shuffle_steps'] > 0 and step % args['shuffle_steps'] == 0:
train_batches.shuffle()
if step % args['eval_steps'] == 0:
dev_score = eval_model(args, trainer, dev_batches, vocab, mwt_dict)
reports = ['Dev score: {:6.3f}'.format(dev_score * 100)]
if step >= args['anneal_after'] and dev_score < prev_dev_score:
reports += ['lr: {:.6f} -> {:.6f}'.format(lr, lr * args['anneal'])]
lr *= args['anneal']
trainer.change_lr(lr)
prev_dev_score = dev_score
if dev_score > best_dev_score:
reports += ['New best dev score!']
best_dev_score = dev_score
best_dev_step = step
trainer.save(args['save_name'])
elif best_dev_step > 0 and step - best_dev_step > args['max_steps_before_stop']:
reports += ['Stopping training after {} steps with no improvement'.format(step - best_dev_step)]
logger.info('\t'.join(reports))
break
logger.info('\t'.join(reports))
if best_dev_step > -1:
logger.info('Best dev score={} at step {}'.format(best_dev_score, best_dev_step))
else:
logger.info('Dev set never evaluated. Saving final model')
trainer.save(args['save_name'])
def evaluate(args):
mwt_dict = load_mwt_dict(args['mwt_json_file'])
use_cuda = args['cuda'] and not args['cpu']
trainer = Trainer(model_file=args['load_name'] or args['save_name'], use_cuda=use_cuda)
loaded_args, vocab = trainer.args, trainer.vocab
for k in loaded_args:
if not k.endswith('_file') and k not in ['cuda', 'mode', 'save_dir', 'load_name', 'save_name']:
args[k] = loaded_args[k]
eval_input_files = {
'json': args['json_file'],
'txt': args['txt_file'],
'label': args['label_file']
}
batches = DataLoader(args, input_files=eval_input_files, vocab=vocab, evaluation=True)
oov_count, N, _, _ = output_predictions(args['conll_file'], trainer, batches, vocab, mwt_dict, args['max_seqlen'])
logger.info("OOV rate: {:6.3f}% ({:6d}/{:6d})".format(oov_count / N * 100, oov_count, N))
if __name__ == '__main__':
main()
|
