-
Notifications
You must be signed in to change notification settings - Fork 15
/
preprocess.py
163 lines (131 loc) · 5.82 KB
/
preprocess.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
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Pre-process Data / features files and build vocabulary
"""
import codecs
import glob
import sys
import gc
import torch
from functools import partial
from onmt.utils.logging import init_logger, logger
from onmt.utils.misc import split_corpus
import onmt.inputters as inputters
import onmt.opts as opts
from onmt.utils.parse import ArgumentParser
def check_existing_pt_files(opt):
""" Check if there are existing .pt files to avoid overwriting them """
pattern = opt.save_data + '.{}*.pt'
for t in ['train', 'valid', 'vocab']:
path = pattern.format(t)
if glob.glob(path):
sys.stderr.write("Please backup existing pt files: %s, "
"to avoid overwriting them!\n" % path)
sys.exit(1)
def build_save_dataset(corpus_type, fields, src_reader, history_reader, ans_reader, tgt_reader, opt):
assert corpus_type in ['train', 'valid']
if corpus_type == 'train':
src = opt.train_src
history = opt.train_history
ans = opt.train_ans
tgt = opt.train_tgt
else:
src = opt.valid_src
history = opt.valid_history
ans = opt.valid_ans
tgt = opt.valid_tgt
logger.info("Reading source and target files: %s %s %s %s." % (src, history, ans, tgt))
src_shards = split_corpus(src, opt.shard_size)
history_shards = split_corpus(history, opt.shard_size)
ans_shards = split_corpus(ans, opt.shard_size)
tgt_shards = split_corpus(tgt, opt.shard_size)
shard_pairs = zip(src_shards, history_shards, ans_shards, tgt_shards)
dataset_paths = []
if (corpus_type == "train" or opt.filter_valid) and tgt is not None:
filter_pred = partial(
inputters.filter_example, use_src_len=opt.data_type == "text", use_history_len=False,
max_src_len=opt.src_seq_length, max_history_len=-1, max_tgt_len=opt.tgt_seq_length)
else:
filter_pred = None
logger.info("filter_pred is not used:{}".format(filter_pred))
for i, (src_shard, history_shard, ans_shard, tgt_shard) in enumerate(shard_pairs):
assert len(src_shard) == len(tgt_shard) and len(src_shard) == len(history_shard) and len(src_shard) == len(ans_shard)
logger.info("Building shard %d." % i)
dataset = inputters.Dataset(
fields,
readers=[src_reader, history_reader, ans_reader, tgt_reader] if tgt_reader else [src_reader, history_reader, ans_reader],
data=([("src", src_shard), ("history", history_shard), ("ans", ans_shard), ("tgt", tgt_shard)]
if tgt_reader else [("src", src_shard), ("history", history_shard), ("ans", ans_shard)]),
dirs=[opt.src_dir, opt.src_dir, opt.src_dir, None] if tgt_reader else [opt.src_dir, opt.src_dir, opt.src_dir],
sort_key=inputters.str2sortkey[opt.data_type],
filter_pred=None
)
data_path = "{:s}.{:s}.{:d}.pt".format(opt.save_data, corpus_type, i)
dataset_paths.append(data_path)
logger.info(" * saving %sth %s data shard to %s."
% (i, corpus_type, data_path))
dataset.save(data_path)
del dataset.examples
gc.collect()
del dataset
gc.collect()
return dataset_paths
def build_save_vocab(train_dataset, fields, opt):
fields = inputters.build_vocab(
train_dataset, fields, opt.data_type, opt.share_vocab,
opt.src_vocab, opt.src_vocab_size, opt.src_words_min_frequency,
opt.tgt_vocab, opt.tgt_vocab_size, opt.tgt_words_min_frequency,
vocab_size_multiple=opt.vocab_size_multiple
)
vocab_path = opt.save_data + '.vocab.pt'
torch.save(fields, vocab_path)
def count_features(path):
"""
path: location of a corpus file with whitespace-delimited tokens and
￨-delimited features within the token
returns: the number of features in the dataset
"""
with codecs.open(path, "r", "utf-8") as f:
first_tok = f.readline().split(None, 1)[0]
return len(first_tok.split(u"￨")) - 1
def main(opt):
ArgumentParser.validate_preprocess_args(opt)
torch.manual_seed(opt.seed)
check_existing_pt_files(opt)
init_logger(opt.log_file)
logger.info("Extracting features...")
src_nfeats = count_features(opt.train_src) if opt.data_type == 'text' \
else 0
tgt_nfeats = count_features(opt.train_tgt) # tgt always text so far
logger.info(" * number of source features: %d." % src_nfeats)
logger.info(" * number of target features: %d." % tgt_nfeats)
logger.info("Building `Fields` object...")
fields = inputters.get_fields(
opt.data_type,
src_nfeats,
tgt_nfeats,
dynamic_dict=opt.dynamic_dict,
src_truncate=opt.src_seq_length_trunc,
tgt_truncate=opt.tgt_seq_length_trunc)
src_reader = inputters.str2reader[opt.data_type].from_opt(opt)
history_reader = inputters.str2reader[opt.data_type].from_opt(opt)
ans_reader = inputters.str2reader[opt.data_type].from_opt(opt)
tgt_reader = inputters.str2reader["text"].from_opt(opt)
logger.info("Building & saving training data...")
train_dataset_files = build_save_dataset(
'train', fields, src_reader, history_reader, ans_reader, tgt_reader, opt)
if opt.valid_src and opt.valid_tgt and opt.valid_history and opt.valid_ans:
logger.info("Building & saving validation data...")
build_save_dataset('valid', fields, src_reader, history_reader, ans_reader, tgt_reader, opt)
logger.info("Building & saving vocabulary...")
build_save_vocab(train_dataset_files, fields, opt)
def _get_parser():
parser = ArgumentParser(description='preprocess.py')
opts.config_opts(parser)
opts.preprocess_opts(parser)
return parser
if __name__ == "__main__":
parser = _get_parser()
opt = parser.parse_args()
main(opt)