-
Notifications
You must be signed in to change notification settings - Fork 45
/
dataset.py
341 lines (298 loc) · 10.6 KB
/
dataset.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
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
"""
Module for dataset creation.
Usage:
python3 dataset.py --command create_dataset --input data/train-v1.1.json --output data/train.bin,data/validation.bin --split 0.8,0.2
python3 dataset.py --command create_vocab --input data/train-v1.1.json --output data/vocab
Modified ver. of https://github.com/tensorflow/models/blob/master/textsum/data.py
"""
import glob
import json
import struct
from random import shuffle
import tensorflow as tf
from tensorflow.core.example import example_pb2
from spacy.en import English
nlp = English()
FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_string('command', 'create_dataset',
'Either create_vocab or create_dataset.'
'Specify FLAGS.in_directories accordingly.')
tf.app.flags.DEFINE_string('input', '', 'path to input data')
tf.app.flags.DEFINE_string('output', '', 'comma separated paths to files')
tf.app.flags.DEFINE_string('split', '', 'comma separated fractions of training/validation')
# special tokens
PARAGRAPH_START = '<p>'
PARAGRAPH_END = '</p>'
SENTENCE_START = '<s>'
SENTENCE_END = '</s>'
UNKNOWN_TOKEN = '<unk>'
PAD_TOKEN = '<pad>'
# special tokens for OOVs
WORD_BEGIN = '<b>'
WORD_CONTINUE = '<c>'
WORD_END = '<e>'
class Vocab:
"""Vocab class for mapping token and ids."""
def __init__(self, file_path, max_size):
self._token_to_id = {}
self._id_to_token = {}
self._size = 0
with open(file_path, 'rt', encoding='utf-8') as f:
for line in f:
tokens = line.split()
# take care of white spaces
if len(tokens) == 1:
count = tokens[0]
idx = line.index(count)
t = line[:idx-1]
tokens = (t, count)
if len(tokens) != 2:
continue
# duplicates
if tokens[0] in self._token_to_id:
continue
self._size += 1
if self._size > max_size:
tf.logging.warn('Warning! Too many tokens: >%d\n' % max_size)
break
self._token_to_id[tokens[0]] = self._size
self._id_to_token[self._size] = tokens[0]
def __len__(self):
return self._size
def tokenToId(self, token):
if token not in self._token_to_id:
tf.logging.warn('id not found for token: %s\n' % token)
return self._token_to_id[UNKNOWN_TOKEN]
return self._token_to_id[token]
def idToToken(self, _id):
if _id not in self._id_to_token:
tf.logging.warn('token not found for id: %d\n' % _id)
return UNKNOWN_TOKEN
return self._id_to_token[_id]
def create_vocab(input_file, output_file, max_size=200000):
"""Generates vocab from input_file.
Args:
input_file: input file path
output_file: output file path
max_size: size of Vocabulary
"""
from collections import Counter
counter = Counter()
with open(input_file, 'r', encoding='utf-8') as data_file:
parsed_file = json.load(data_file)
data = parsed_file['data']
for datum in data:
for paragraph in datum['paragraphs']:
context = nlp(paragraph['context'].lower())
counter.update(context.text)
counter.update(map(lambda c: c.text, context))
for qas in paragraph['qas']:
question = nlp(qas['question'].lower())
counter.update(question.text)
counter.update(map(lambda c: c.text, question))
with open(output_file, 'wt') as f:
# reserve for special tokens
f.write('<s> 0\n')
f.write('</s> 0\n')
f.write('<unk> 0\n')
f.write('<pad> 0\n')
f.write('<b> 0\n')
f.write('<c> 0\n')
f.write('<e> 0\n')
for token, count in counter.most_common(max_size-7):
f.write(token + ' ' + str(count) + '\n')
def create_dataset(input_file, output_files, split_fractions):
"""Generates train/validation files from input_file.
Args:
input_file: input file path
output_file: output file path
split_fractions: train/validation split fractions
"""
import struct
from random import shuffle
from nltk.tokenize import sent_tokenize
from tensorflow.core.example import example_pb2
with open(input_file, 'r') as data_file:
parsed_file = json.load(data_file)
data = parsed_file['data']
len_data = len(data)
indices = [int(len_data*(1-split)) for split in split_fractions]
indices.insert(0, 0)
# shuffle data by topic
shuffle(data)
for i in range(1, len(indices)):
subset = data[indices[i-1]:indices[i]]
with open(output_files[i-1], 'wb') as writer:
for datum in subset:
for paragraph in datum['paragraphs']:
context = nlp(paragraph['context']).text
sentences = sent_tokenize(context)
context = '<p>' + ' '.join(['<s>' + sentence + '</s>' for sentence in sentences]) + '</p>'
context = context.encode('utf-8')
qas = paragraph['qas']
for qa in qas:
question = nlp(qa['question']).text
answer = nlp(qa['answers'][0]['text']).text # just select best one
sentences = sent_tokenize(question)
question = '<p>' + ' '.join(['<s>' + sentence + '</s>' for sentence in sentences]) + '</p>'
question = question.encode('utf-8')
sentences = sent_tokenize(answer)
answer = '<p>' + ' '.join(['<s>' + sentence + '</s>' for sentence in sentences]) + '</p>'
answer = answer.encode('utf-8')
tf_example = example_pb2.Example()
tf_example.features.feature['context'].bytes_list.value.extend([context])
tf_example.features.feature['question'].bytes_list.value.extend([question])
tf_example.features.feature['answer'].bytes_list.value.extend([answer])
tf_example_str = tf_example.SerializeToString()
str_len = len(tf_example_str)
writer.write(struct.pack('q', str_len))
writer.write(struct.pack('%ds' % str_len, tf_example_str))
def snippet_gen(text, start_tok, end_tok, inclusive=False):
"""Generates consecutive snippets between start and end tokens.
Args:
text: a string
start_tok: a string denoting the start of snippets
end_tok: a string denoting the end of snippets
inclusive: Whether include the tokens in the returned snippets.
Yields:
String snippets
"""
cur = 0
while True:
try:
start_p = text.index(start_tok, cur)
end_p = text.index(end_tok, start_p + 1)
cur = end_p + len(end_tok)
if inclusive:
yield text[start_p:cur]
else:
yield text[start_p+len(start_tok):end_p]
except ValueError as e:
raise StopIteration('no more snippets in text: %s' % e)
def to_sentences(paragraph, include_token=False):
"""Takes tokens of a paragraph and returns list of sentences.
Args:
paragraph: string, text of paragraph
include_token: Whether include the sentence separation tokens result.
Returns:
List of sentence strings.
"""
if not isinstance(paragraph, str):
paragraph = paragraph.decode('utf-8')
s_gen = snippet_gen(paragraph, SENTENCE_START, SENTENCE_END, include_token)
return [s for s in s_gen]
def pad(ids, pad_id, length):
"""Pad or trim list to len length.
Args:
ids: list of ints to pad
pad_id: what to pad with
length: length to pad or trim to
Returns:
ids trimmed or padded with pad_id
"""
assert pad_id is not None
assert length is not None
if len(ids) < length:
a = [pad_id] * (length - len(ids))
return ids + a
else:
return ids[:length]
def tokens_to_ids(text, vocab, pad_len=None, pad_id=None):
"""Get ids corresponding to tokens in text.
Assumes tokens separated by space.
Args:
text: a string
vocab: TextVocabularyFile object
pad_len: int, length to pad to
pad_id: int, token id for pad symbol
Returns:
A list of ints representing token ids.
"""
ids = []
b = vocab.tokenToId(WORD_BEGIN)
c = vocab.tokenToId(WORD_CONTINUE)
e = vocab.tokenToId(WORD_END)
unk = vocab.tokenToId(UNKNOWN_TOKEN)
token_iterator = map(lambda x: x.text, nlp(text.lower()))
for token in token_iterator:
_id = vocab.tokenToId(token)
if _id == unk: # w is OOV
ids.append(b)
for character in token:
ids.append(c)
ids.append(vocab.tokenToId(character))
ids.append(e)
else: # w is present in vocab
ids.append(_id)
if pad_len is not None:
return pad(ids, pad_id, pad_len)
return ids
def ids_to_tokens(ids_list, vocab):
"""Get tokens from ids.
Args:
ids_list: list of int32
vocab: TextVocabulary object
Returns:
List of tokens corresponding to ids.
"""
assert isinstance(ids_list, list), '%s is not a list' % ids_list
answer = []
tmp = ''
# iterate throught each id and recover any OOVs
for _id in ids_list:
token = vocab.idToToken(_id)
if token == PAD_TOKEN:
token = ''
if token == WORD_BEGIN:
tmp += token
elif token == WORD_END:
tmp = ''.join(tmp.split(WORD_CONTINUE))
answer.append(tmp[1:])
tmp = ''
elif len(tmp) > 0:
tmp += token
else:
answer.append(token)
return answer
def tf_Examples(data_path, num_epochs=None):
"""Generates tf.Examples from path of data files.
Binary data format: <length><blob>. <length> represents the byte size
of <blob>. <blob> is serialized tf.Example proto. The tf.Example contains
the tokenized article text and summary.
Args:
data_path: path to tf.Example data files.
num_epochs: Number of times to go through the data. None means infinite.
Yields:
Deserialized tf.Example.
If there are multiple files specified, they accessed in a random order.
"""
epoch = 0
while True:
if num_epochs is not None and epoch >= num_epochs:
break
filelist = glob.glob(data_path)
assert filelist, 'Empty filelist.'
shuffle(filelist)
for f in filelist:
reader = open(f, 'rb')
while True:
len_bytes = reader.read(8)
if not len_bytes: break
str_len = struct.unpack('q', len_bytes)[0]
example_str = struct.unpack('%ds' % str_len, reader.read(str_len))[0]
yield example_pb2.Example.FromString(example_str)
epoch += 1
def main(unused_argv):
assert FLAGS.command and FLAGS.input and FLAGS.output
output_files = FLAGS.output.split(',')
input_file = FLAGS.input
if FLAGS.command == 'create_dataset':
assert FLAGS.split
split_fractions = [float(s) for s in FLAGS.split.split(',')]
assert len(output_files) == len(split_fractions)
create_dataset(input_file, output_files, split_fractions)
elif FLAGS.command == 'create_vocab':
assert len(output_files) == 1
create_vocab(input_file, output_files[0])
if __name__ == '__main__':
tf.app.run()