-
Notifications
You must be signed in to change notification settings - Fork 8
/
preprocess.py
254 lines (222 loc) · 13.4 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
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
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
import itertools
import json
import multiprocessing as mp
import pickle
from argparse import ArgumentParser
from collections import defaultdict
import tensorflow as tf
MAX_CONTEXTS = 0
MAX_INTERNAL_PATHS = 0
MAX_PATH_LENGTH = 0
MAX_RELATIVE_PATH_LENGTH = 0
MAX_EXAMPLES_IN_SHARD = 500000
def save_dictionaries(dataset_name, subtok_to_count, node_to_count, max_contexts, max_internal_paths,
max_path_length, max_path_width, max_relative_path_length, max_child_id):
save_dict_file_path = '{}.dict.cg'.format(dataset_name)
with open(save_dict_file_path, 'wb') as file:
pickle.dump(subtok_to_count, file)
pickle.dump(node_to_count, file)
pickle.dump(max_contexts, file)
pickle.dump(max_internal_paths, file)
pickle.dump(max_path_length, file)
pickle.dump(max_path_width, file)
pickle.dump(max_relative_path_length, file)
pickle.dump(max_child_id, file)
print('Dictionaries saved to: {}'.format(save_dict_file_path))
def make_example_from_line(line):
obj = json.loads(line)
return make_example(obj)
def make_example_and_histograms(line):
local_node_to_count = defaultdict(int)
local_subtoken_to_count = defaultdict(int)
local_total_paths, local_total_examples = 0, 0
obj = json.loads(line)
targets = obj['targets']
is_token_flags = obj['is_token']
paths = obj['head_paths'] + obj['relative_paths']
tokens = set()
for target_val, is_tok in zip(targets, is_token_flags):
if is_tok == 1:
if ',' in target_val:
values = target_val.split(',')
for v in values:
local_subtoken_to_count[v] += 1
else:
local_subtoken_to_count[target_val] += 1
tokens.add(target_val)
else:
local_node_to_count[target_val] += 1
for path in paths:
for subtok in path['sources']:
local_subtoken_to_count[subtok] += 1
for node_with_child in path['nodes']:
node_with_child = node_with_child['node']
if node_with_child[0] in tokens:
continue
node = node_with_child[0].rstrip('_INV')
local_node_to_count[node] += 1
if not node.startswith('@'):
local_node_to_count[node + '_INV'] += 1
local_total_paths += 1
local_total_examples += len(targets)
ex = make_example(obj)
return ex, local_node_to_count, local_subtoken_to_count, local_total_paths, local_total_examples
def make_example(obj):
ex = tf.train.SequenceExample()
ex.context.feature['num_targets'].int64_list.value.append(obj['num_targets'])
ex.context.feature['num_nodes'].int64_list.value.append(obj['num_nodes'])
ex.context.feature['data_num_contexts'].int64_list.value.append(MAX_CONTEXTS)
ex.context.feature['max_internal_paths'].int64_list.value.append(MAX_INTERNAL_PATHS)
ex.context.feature['head_target_child_id'].int64_list.value.append(obj['head_child_id'])
ex.context.feature['linearized_tree'].bytes_list.value.append(obj['linearized_tree'].encode())
ex.context.feature['filepath'].bytes_list.value.append(obj['filepath'].encode())
ex.context.feature['line'].int64_list.value.append(obj['line'])
targets = ex.feature_lists.feature_list['targets']
is_token = ex.feature_lists.feature_list['is_token']
target_child_id = ex.feature_lists.feature_list['target_child_id']
relative_path_nodes = ex.feature_lists.feature_list['relative_path_nodes']
relative_child_ids = ex.feature_lists.feature_list['relative_path_child_ids']
internal_path_nodes = ex.feature_lists.feature_list['internal_path_nodes']
internal_path_nodes_sources = ex.feature_lists.feature_list['internal_paths_sources']
internal_path_nodes_child_ids = ex.feature_lists.feature_list['internal_paths_child_ids']
head_paths_sources = ex.feature_lists.feature_list['head_paths_sources']
head_paths_tokens = ex.feature_lists.feature_list['head_paths_tokens']
head_paths_nodes = ex.feature_lists.feature_list['head_paths_nodes']
head_paths_child_ids = ex.feature_lists.feature_list['head_paths_child_ids']
head_root_nodes = ex.feature_lists.feature_list['head_root_nodes']
head_root_child_ids = ex.feature_lists.feature_list['head_root_child_ids']
for target in obj['targets']:
targets.feature.add().bytes_list.value.append(target.encode())
for is_tok in obj['is_token']:
is_token.feature.add().int64_list.value.append(is_tok)
for tgt_child in obj['target_child_id']:
target_child_id.feature.add().int64_list.value.append(tgt_child)
for rps in obj['relative_paths']:
node_child_pairs = rps['nodes'][-MAX_RELATIVE_PATH_LENGTH:]
relative_path_nodes.feature.add().bytes_list.value.extend(nc['node'][0].encode() for nc in node_child_pairs)
relative_child_ids.feature.add().int64_list.value.extend(int(nc['node'][1]) for nc in node_child_pairs)
for paths_for_target in obj['internal_paths']:
for i, path in enumerate(paths_for_target):
if i >= MAX_INTERNAL_PATHS:
break
internal_path_nodes_sources.feature.add().bytes_list.value.extend(subtok.encode() for subtok in path['sources'])
nodes_with_childs = path['nodes'][-MAX_PATH_LENGTH:]
node_child_pairs = [nc['node'] for nc in nodes_with_childs]
internal_path_nodes.feature.add().bytes_list.value.extend(nc[0].encode() for nc in node_child_pairs)
internal_path_nodes_child_ids.feature.add().int64_list.value.extend(int(nc[1]) for nc in node_child_pairs)
for j in range(MAX_INTERNAL_PATHS - min(MAX_INTERNAL_PATHS, len(paths_for_target))):
internal_path_nodes_sources.feature.add()
internal_path_nodes.feature.add()
internal_path_nodes_child_ids.feature.add()
sorted_head_paths = sorted(obj['head_paths'], key=lambda x: len(x['nodes']))
for head_path in sorted_head_paths[:MAX_CONTEXTS]:
node_child_pairs = head_path['nodes'][:MAX_PATH_LENGTH]
head_paths_sources.feature.add().bytes_list.value.extend(subtok.encode() for subtok in head_path['sources'])
head_paths_tokens.feature.add().bytes_list.value.append(','.join(head_path['sources']).encode())
head_paths_nodes.feature.add().bytes_list.value.extend(nc['node'][0].encode() for nc in node_child_pairs)
head_paths_child_ids.feature.add().int64_list.value.extend(int(nc['node'][1]) for nc in node_child_pairs)
for node in obj['head_root_path']['nodes'][-MAX_PATH_LENGTH:]:
node = node['node']
head_root_nodes.feature.add().bytes_list.value.append(node[0].encode())
head_root_child_ids.feature.add().int64_list.value.append(node[1])
return ex.SerializeToString()
def process_file(file_path, data_file_role, dataset_name, max_contexts, max_internal_paths, max_path_length,
max_relative_path_length, collect_histograms=False):
# Currently we take max contexts both from this script and from the json.
# When moving to joint paths, we should pad here and take max_contexts from the arguments and not the json
total_paths = 0
total_examples = 0
subtoken_to_count = defaultdict(int)
node_to_count = defaultdict(int)
global MAX_CONTEXTS, MAX_INTERNAL_PATHS, MAX_PATH_LENGTH, MAX_RELATIVE_PATH_LENGTH
MAX_CONTEXTS = max_contexts
MAX_INTERNAL_PATHS = max_internal_paths
MAX_PATH_LENGTH = max_path_length
MAX_RELATIVE_PATH_LENGTH = max_relative_path_length
with open(file_path, 'r') as file:
current_shard_number = 0
writer = create_writer(current_shard_number, data_file_role, dataset_name)
if collect_histograms:
with mp.Pool(64) as pool:
examples_with_histograms = pool.imap_unordered(make_example_and_histograms, file, chunksize=100)
#examples_with_histograms = [make_example_and_histograms(line) for line in file]
for i, (ex, local_node_to_count, local_subtoken_to_count, local_total_paths,
local_total_examples) in enumerate(examples_with_histograms):
for key, val in local_node_to_count.items():
node_to_count[key] += val
for key, val in local_subtoken_to_count.items():
subtoken_to_count[key] += val
total_paths += local_total_paths
total_examples += local_total_examples
if (i+1) % MAX_EXAMPLES_IN_SHARD == 0:
current_shard_number += 1
writer.close()
writer = create_writer(current_shard_number, data_file_role, dataset_name)
writer.write(ex)
else:
with mp.Pool(64) as pool:
serialized_examples = pool.imap_unordered(make_example_from_line, file, chunksize=100)
#serialized_examples = [make_example_from_line(line) for line in file]
for i, ex in enumerate(serialized_examples):
if (i+1) % MAX_EXAMPLES_IN_SHARD == 0:
current_shard_number += 1
writer.close()
writer = create_writer(current_shard_number, data_file_role, dataset_name)
writer.write(ex)
writer.close()
print('File: ' + file_path)
if collect_histograms:
print('Average total contexts: ' + str(float(total_paths) / total_examples))
print('Total examples: ' + str(total_examples))
return total_examples, subtoken_to_count, node_to_count
def create_writer(current_shard_number, data_file_role, dataset_name):
output_path = '{}.tfrecord.{}.{}.cg'.format(dataset_name, data_file_role, current_shard_number)
writer = tf.io.TFRecordWriter(output_path, options=tf.io.TFRecordCompressionType.GZIP)
return writer
if __name__ == '__main__':
parser = ArgumentParser()
parser.add_argument("-trd", "--train_data", dest="train_data_path",
help="path to training data file", required=True)
parser.add_argument("-ted", "--test_data", dest="test_data_path",
help="path to test data file", required=True)
parser.add_argument("-vd", "--val_data", dest="val_data_path",
help="path to validation data file", required=True)
parser.add_argument("-mc", "--max_contexts", dest="max_contexts", default=200,
help="number of max contexts to keep", required=False)
parser.add_argument("--max_internal_paths", dest="max_internal_paths",
help="number of max internal paths to keep", required=True)
parser.add_argument("-mp", "--max_path_length", dest="max_path_length", default=12,
required=False)
parser.add_argument("--max_relative_path_length", dest="max_relative_path_length", default=6,
required=False)
parser.add_argument("-mw", "--max_path_width", dest="max_path_width", default=3,
required=False)
parser.add_argument("--max_child_id", dest="max_child_id", default=5,
help="number of max nodes to keep", required=False)
parser.add_argument("-svs", "--subtoken_vocab_size", dest="subtoken_vocab_size", default=186277,
help="Max number of source subtokens to keep in the vocabulary", required=False)
parser.add_argument("-o", "--output_name", dest="output_name",
help="output name - the base name for the created dataset", metavar="FILE", required=True,
default='data')
args = parser.parse_args()
train_data_path = args.train_data_path
test_data_path = args.test_data_path
val_data_path = args.val_data_path
num_examples, subtoken_to_count, node_to_count = process_file(file_path=train_data_path, data_file_role='train',
dataset_name=args.output_name,
max_contexts=int(args.max_contexts),
max_internal_paths=int(args.max_internal_paths),
max_path_length=int(args.max_path_length),
max_relative_path_length=int(args.max_relative_path_length),
collect_histograms=True)
for data_file_path, data_role in zip([test_data_path, val_data_path], ['test', 'val']):
process_file(file_path=data_file_path, data_file_role=data_role, dataset_name=args.output_name,
max_contexts=int(args.max_contexts), max_internal_paths=int(args.max_internal_paths),
max_path_length=int(args.max_path_length), max_relative_path_length=int(args.max_relative_path_length),
collect_histograms=False)
save_dictionaries(dataset_name=args.output_name, subtok_to_count=subtoken_to_count,
node_to_count=node_to_count,
max_contexts=int(args.max_contexts), max_internal_paths=int(args.max_internal_paths),
max_path_length=int(args.max_path_length), max_path_width=int(args.max_path_width),
max_relative_path_length=int(args.max_relative_path_length),
max_child_id=int(args.max_child_id))