-
Notifications
You must be signed in to change notification settings - Fork 9
/
preprocess_captions.py
200 lines (169 loc) · 5.66 KB
/
preprocess_captions.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
# this file is most borrowd from an open-access blog.
# I am not sure who owns the copyright. --lukun199
# Formatting the raw data take a few minutes.
import string
import re, os
from pickle import dump
from unicodedata import normalize
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--data_root", type=str, default='')
args = parser.parse_args()
# load doc into memory
def load_doc(filename):
# open the file as read only
file = open(filename, mode='rt', encoding='utf-8')
# read all text
text = file.read()
# close the file
file.close()
return text
# split a loaded document into sentences
def to_sentences(doc):
return doc.strip().split('\n')
# shortest and longest sentence lengths
def sentence_lengths(sentences):
lengths = [len(s.split()) for s in sentences]
return min(lengths), max(lengths)
# clean a list of lines
def clean_lines(lines):
cleaned = list()
# prepare regex for char filtering
re_print = re.compile('[^%s]' % re.escape(string.printable))
# prepare translation table for removing punctuation
table = str.maketrans('', '', string.punctuation)
for line in lines:
# normalize unicode characters
line = normalize('NFD', line).encode('ascii', 'ignore')
line = line.decode('UTF-8')
# tokenize on white space
line = line.split()
# convert to lower case
line = [word.lower() for word in line]
# remove punctuation from each token
line = [word.translate(table) for word in line]
# remove non-printable chars form each token
line = [re_print.sub('', w) for w in line]
# remove tokens with numbers in them
line = [word for word in line if word.isalpha()]
# store as string
cleaned.append(' '.join(line))
return cleaned
# load doc into memory
def load_doc(filename):
# open the file as read only
file = open(filename, mode='rt', encoding='utf-8')
# read all text
text = file.read()
# close the file
file.close()
return text
# split a loaded document into sentences
def to_sentences(doc):
return doc.strip().split('\n')
# clean a list of lines
def clean_lines(lines):
cleaned = list()
# prepare regex for char filtering
re_print = re.compile('[^%s]' % re.escape(string.printable))
# prepare translation table for removing punctuation
table = str.maketrans('', '', string.punctuation)
for idx, line in enumerate(lines):
# normalize unicode characters
line = normalize('NFD', line).encode('ascii', 'ignore')
line = line.decode('UTF-8')
# tokenize on white space
line = line.split()
# convert to lower case
line = [word.lower() for word in line]
# remove punctuation from each token
line = [word.translate(table) for word in line]
# remove non-printable chars form each token
line = [re_print.sub('', w) for w in line]
# remove tokens with numbers in them
line = [word for word in line if word.isalpha()]
# store as string
cleaned.append(' '.join(line))
if (idx % 100000) == 0:
print(idx)
return cleaned
# save a list of clean sentences to file
def save_clean_sentences(sentences, filename):
dump(sentences, open(filename, 'wb'))
print('Saved: %s' % filename)
# load English data
root = args.data_root
filename = 'europarl-v7.fr-en.en'
doc = load_doc(os.path.join(root, filename))
sentences = to_sentences(doc)
minlen, maxlen = sentence_lengths(sentences)
print('English data: sentences=%d, min=%d, max=%d' % (len(sentences), minlen, maxlen))
# clean dataset
sentences = clean_lines(sentences)
save_clean_sentences(sentences, 'english.pkl')
# spot check
for i in range(10):
print(sentences[i])
from pickle import load
from collections import Counter
# load a clean dataset
def load_clean_sentences(filename):
return load(open(filename, 'rb'))
# save a list of clean sentences to file
def save_clean_sentences(sentences, filename):
dump(sentences, open(filename, 'wb'))
print('Saved: %s' % filename)
# create a frequency table for all words
def to_vocab(lines):
vocab = Counter()
for line in lines:
tokens = line.split()
vocab.update(tokens)
return vocab
# remove all words with a frequency below a threshold
def trim_vocab(vocab, min_occurance):
tokens = [k for k, c in vocab.items() if c >= min_occurance]
return set(tokens)
# mark all OOV with "unk" for all lines
def update_dataset(lines, vocab):
new_lines = list()
for line in lines:
new_tokens = list()
for token in line.split():
if token in vocab:
new_tokens.append(token)
else:
new_tokens.append('unk')
new_line = ' '.join(new_tokens)
new_lines.append(new_line)
return new_lines
def clip_length(lines):
ss = []
for x in lines:
if 20>=len(x.split(' '))>=3 :
ss.append(x)
return ss
# load English dataset
filename = 'english.pkl'
lines = load_clean_sentences(filename)
# satisfy the length limit
lines = clip_length(lines)
# calculate vocabulary
vocab = to_vocab(lines)
print('English Vocabulary: %d' % len(vocab))
# reduce vocabulary
vocab = trim_vocab(vocab, min_occurance=5)
print('New English Vocabulary: %d' % len(vocab))
# mark out of vocabulary words
lines = update_dataset(lines, vocab)
# save updated dataset
filename = 'english_vocab.pkl'
to_save = {}
to_save['sent_str'] = lines
to_save['voc'] = {word:i+1 for i,word in enumerate(vocab)}
to_save['voc']['unk'] = 0
to_save['len_range'] = (3,20)
save_clean_sentences(to_save, filename)
# spot check
for i in range(10):
print(lines[i])