/
make_datafiles_train.py
209 lines (170 loc) · 7.96 KB
/
make_datafiles_train.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
import sys
import os
import hashlib
import struct
import subprocess
import collections
import tensorflow as tf
from tensorflow.core.example import example_pb2
dm_single_close_quote = u'\u2019' # unicode
dm_double_close_quote = u'\u201d'
END_TOKENS = ['.', '!', '?', '...', "'", "`", '"', dm_single_close_quote, dm_double_close_quote, ")"] # acceptable ways to end a sentence
# We use these to separate the summary sentences in the .bin datafiles
SENTENCE_START = '<s>'
SENTENCE_END = '</s>'
CHUNK_SIZE = 1000 # num examples per chunk, for the chunked data
VOCAB_SIZE = 200000
def chunk_file(set_name):
# in_file = 'finished_files/%s.bin' % set_name
in_file = os.path.join(finished_files_dir, '%s.bin' % set_name)
reader = open(in_file, "rb")
chunk = 0
finished = False
while not finished:
chunk_fname = os.path.join(chunks_dir, '%s_%03d.bin' % (set_name, chunk)) # new chunk
with open(chunk_fname, 'wb') as writer:
for _ in range(CHUNK_SIZE):
len_bytes = reader.read(8)
if not len_bytes:
finished = True
break
str_len = struct.unpack('q', len_bytes)[0]
example_str = struct.unpack('%ds' % str_len, reader.read(str_len))[0]
writer.write(struct.pack('q', str_len))
writer.write(struct.pack('%ds' % str_len, example_str))
chunk += 1
def chunk_all():
# Make a dir to hold the chunks
if not os.path.isdir(chunks_dir):
os.mkdir(chunks_dir)
# Chunk the data
# for set_name in ['train', 'val', 'test']:
for set_name in ['train']:
print("Splitting %s data into chunks..." % set_name)
chunk_file(set_name)
print("Saved chunked data in %s" % chunks_dir)
def tokenize_stories(stories_dir, tokenized_stories_dir):
"""Maps a whole directory of .story files to a tokenized version using Stanford CoreNLP Tokenizer"""
print("Preparing to tokenize %s to %s..." % (stories_dir, tokenized_stories_dir))
stories = os.listdir(stories_dir)
# make IO list file
print("Making list of files to tokenize...")
with open("mapping.txt", "w") as f:
for s in stories:
f.write("%s \t %s\n" % (os.path.join(stories_dir, s), os.path.join(tokenized_stories_dir, s)))
command = ['java', '-cp', 'stanford-corenlp-full-2018-02-27/stanford-corenlp-3.9.1.jar', 'edu.stanford.nlp.process.PTBTokenizer', '-ioFileList', '-preserveLines', 'mapping.txt']
print("Tokenizing %i files in %s and saving in %s..." % (len(stories), stories_dir, tokenized_stories_dir))
subprocess.call(command)
print("Stanford CoreNLP Tokenizer has finished.")
os.remove("mapping.txt")
# Check that the tokenized stories directory contains the same number of files as the original directory
num_orig = len(os.listdir(stories_dir))
num_tokenized = len(os.listdir(tokenized_stories_dir))
if num_orig != num_tokenized:
raise Exception("The tokenized stories directory %s contains %i files, but it should contain the same number as %s (which has %i files). Was there an error during tokenization?" % (tokenized_stories_dir, num_tokenized, stories_dir, num_orig))
print("Successfully finished tokenizing %s to %s.\n" % (stories_dir, tokenized_stories_dir))
def read_text_file(text_file):
lines = []
with open(text_file, "r") as f:
for line in f:
lines.append(line.strip())
return lines
def fix_missing_period(line):
"""Adds a period to a line that is missing a period"""
if "@highlight" in line: return line
if line=="": return line
if line[-1] in END_TOKENS: return line
# print line[-1]
return line + " ."
def get_art_abs(story_file):
lines = read_text_file(story_file)
# Lowercase everything
lines = [line.lower() for line in lines]
# Put periods on the ends of lines that are missing them (this is a problem in the dataset because many image captions don't end in periods; consequently they end up in the body of the article as run-on sentences)
lines = [fix_missing_period(line) for line in lines]
# Separate out article and abstract sentences
article_lines = []
highlights = []
next_is_highlight = False
for idx,line in enumerate(lines):
if line == "":
continue # empty line
elif line.startswith("@highlight"):
next_is_highlight = True
elif next_is_highlight:
highlights.append(line)
else:
article_lines.append(line)
# Make article into a single string
article = ' '.join(article_lines)
# Make abstract into a signle string, putting <s> and </s> tags around the sentences
abstract = ' '.join(["%s %s %s" % (SENTENCE_START, sent, SENTENCE_END) for sent in highlights])
return article, abstract
def write_to_bin(out_file,makevocab = False):
# story_fnames = [name for name in os.listdir(tokenized_stories_dir) if os.path.isfile(tokenized_stories_dir+'\\'+name) ]
story_fnames = [name for name in os.listdir(tokenized_stories_dir)]
num_stories = len(story_fnames)
# print(story_fnames)
# for idx,s in enumerate(story_fnames):
# print(idx,s)
if makevocab:
vocab_counter = collections.Counter()
with open(out_file, 'wb') as writer:
for idx,s in enumerate(story_fnames):
if idx % 1000 == 0:
print("Writing story %i of %i; %.2f percent done" % (idx, num_stories, float(idx)*100.0/float(num_stories)))
# Look in the tokenized story dirs to find the .story file corresponding to this url
if os.path.isfile(os.path.join(tokenized_stories_dir, s)):
story_file = os.path.join(tokenized_stories_dir, s)
print(story_file)
else:
print('Error: no data.')
# # Get the strings to write to .bin file
article, abstract = get_art_abs(story_file)
# print('article:', article)
# print('abstract:', abstract)
# Write to tf.Example
tf_example = example_pb2.Example()
tf_example.features.feature['article'].bytes_list.value.extend([article.encode()])
tf_example.features.feature['abstract'].bytes_list.value.extend([abstract.encode()])
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))
if makevocab:
art_tokens = article.split(' ')
abs_tokens = abstract.split(' ')
abs_tokens = [t for t in abs_tokens if t not in [SENTENCE_START, SENTENCE_END]] # remove these tags from vocab
tokens = art_tokens + abs_tokens
tokens = [t.strip() for t in tokens] # strip
tokens = [t for t in tokens if t!=""] # remove empty
vocab_counter.update(tokens)
print("Finished writing file %s\n" % out_file)
# write vocab to file
if makevocab:
print("Writing vocab file...")
with open(os.path.join(finished_files_dir, "vocab"), 'w') as writer:
for word, count in vocab_counter.most_common(VOCAB_SIZE):
writer.write(word + ' ' + str(count) + '\n')
print("Finished writing vocab file")
if __name__ == '__main__':
if len(sys.argv) != 3:
print("USAGE: python make_datafiles.py <stories_dir> <out_dir>")
sys.exit()
stories_dir = sys.argv[1]
out_dir = sys.argv[2]
# print(stories_dir)
# tokenized_stories_dir = "data/tokenized_stories_dir"
# finished_files_dir = "data/finished_files"
tokenized_stories_dir = os.path.join(out_dir, "tokenized_stories_dir")
finished_files_dir = os.path.join(out_dir, "finished_files")
chunks_dir = os.path.join(finished_files_dir, "chunked")
# Create some new directories
if not os.path.exists(tokenized_stories_dir): os.makedirs(tokenized_stories_dir)
if not os.path.exists(finished_files_dir): os.makedirs(finished_files_dir)
# Run stanford tokenizer on both stories dirs, outputting to tokenized stories directories
tokenize_stories(stories_dir, tokenized_stories_dir)
# Read the tokenized stories, do a little postprocessing then write to bin files
write_to_bin(os.path.join(finished_files_dir, "train.bin"),makevocab=True)
# # Chunk the data. This splits each of train.bin, val.bin and test.bin into smaller chunks, each containing e.g. 1000 examples, and saves them in finished_files/chunks
chunk_all()