-
Notifications
You must be signed in to change notification settings - Fork 212
/
inference.py
229 lines (190 loc) · 9.62 KB
/
inference.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
import tensorflow as tf
import spacy
import os
import numpy as np
import ujson as json
from func import cudnn_gru, native_gru, dot_attention, summ, ptr_net
from prepro import word_tokenize, convert_idx
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
# Must be consistant with training
char_limit = 16
hidden = 75
char_dim = 8
char_hidden = 100
use_cudnn = True
# File path
target_dir = "data"
save_dir = "log/model"
word_emb_file = os.path.join(target_dir, "word_emb.json")
char_emb_file = os.path.join(target_dir, "char_emb.json")
word2idx_file = os.path.join(target_dir, "word2idx.json")
char2idx_file = os.path.join(target_dir, "char2idx.json")
class InfModel(object):
# Used to zero elements in the probability matrix that correspond to answer
# spans that are longer than the number of tokens specified here.
max_answer_tokens = 15
def __init__(self, word_mat, char_mat):
self.c = tf.placeholder(tf.int32, [1, None])
self.q = tf.placeholder(tf.int32, [1, None])
self.ch = tf.placeholder(tf.int32, [1, None, char_limit])
self.qh = tf.placeholder(tf.int32, [1, None, char_limit])
self.tokens_in_context = tf.placeholder(tf.int64)
self.word_mat = tf.get_variable("word_mat", initializer=tf.constant(
word_mat, dtype=tf.float32), trainable=False)
self.char_mat = tf.get_variable(
"char_mat", initializer=tf.constant(char_mat, dtype=tf.float32))
self.c_mask = tf.cast(self.c, tf.bool)
self.q_mask = tf.cast(self.q, tf.bool)
self.c_len = tf.reduce_sum(tf.cast(self.c_mask, tf.int32), axis=1)
self.q_len = tf.reduce_sum(tf.cast(self.q_mask, tf.int32), axis=1)
self.c_maxlen = tf.reduce_max(self.c_len)
self.q_maxlen = tf.reduce_max(self.q_len)
self.ch_len = tf.reshape(tf.reduce_sum(
tf.cast(tf.cast(self.ch, tf.bool), tf.int32), axis=2), [-1])
self.qh_len = tf.reshape(tf.reduce_sum(
tf.cast(tf.cast(self.qh, tf.bool), tf.int32), axis=2), [-1])
self.ready()
def ready(self):
N, PL, QL, CL, d, dc, dg = \
1, self.c_maxlen, self.q_maxlen, char_limit, hidden, char_dim, \
char_hidden
gru = cudnn_gru if use_cudnn else native_gru
with tf.variable_scope("emb"):
with tf.variable_scope("char"):
ch_emb = tf.reshape(tf.nn.embedding_lookup(
self.char_mat, self.ch), [N * PL, CL, dc])
qh_emb = tf.reshape(tf.nn.embedding_lookup(
self.char_mat, self.qh), [N * QL, CL, dc])
cell_fw = tf.contrib.rnn.GRUCell(dg)
cell_bw = tf.contrib.rnn.GRUCell(dg)
_, (state_fw, state_bw) = tf.nn.bidirectional_dynamic_rnn(
cell_fw, cell_bw, ch_emb, self.ch_len, dtype=tf.float32)
ch_emb = tf.concat([state_fw, state_bw], axis=1)
_, (state_fw, state_bw) = tf.nn.bidirectional_dynamic_rnn(
cell_fw, cell_bw, qh_emb, self.qh_len, dtype=tf.float32)
qh_emb = tf.concat([state_fw, state_bw], axis=1)
qh_emb = tf.reshape(qh_emb, [N, QL, 2 * dg])
ch_emb = tf.reshape(ch_emb, [N, PL, 2 * dg])
with tf.name_scope("word"):
c_emb = tf.nn.embedding_lookup(self.word_mat, self.c)
q_emb = tf.nn.embedding_lookup(self.word_mat, self.q)
c_emb = tf.concat([c_emb, ch_emb], axis=2)
q_emb = tf.concat([q_emb, qh_emb], axis=2)
with tf.variable_scope("encoding"):
rnn = gru(num_layers=3, num_units=d, batch_size=N,
input_size=c_emb.get_shape().as_list()[-1])
c = rnn(c_emb, seq_len=self.c_len)
q = rnn(q_emb, seq_len=self.q_len)
with tf.variable_scope("attention"):
qc_att = dot_attention(c, q, mask=self.q_mask, hidden=d)
rnn = gru(num_layers=1, num_units=d, batch_size=N,
input_size=qc_att.get_shape().as_list()[-1])
att = rnn(qc_att, seq_len=self.c_len)
with tf.variable_scope("match"):
self_att = dot_attention(att, att, mask=self.c_mask, hidden=d)
rnn = gru(num_layers=1, num_units=d, batch_size=N,
input_size=self_att.get_shape().as_list()[-1])
match = rnn(self_att, seq_len=self.c_len)
with tf.variable_scope("pointer"):
init = summ(q[:, :, -2 * d:], d, mask=self.q_mask)
pointer = ptr_net(batch=N, hidden=init.get_shape().as_list()[-1])
logits1, logits2 = pointer(init, match, d, self.c_mask)
with tf.variable_scope("predict"):
outer = tf.matmul(tf.expand_dims(tf.nn.softmax(logits1), axis=2),
tf.expand_dims(tf.nn.softmax(logits2), axis=1))
outer = tf.cond(
self.tokens_in_context < self.max_answer_tokens,
lambda: tf.matrix_band_part(outer, 0, -1),
lambda: tf.matrix_band_part(outer, 0, self.max_answer_tokens))
self.yp1 = tf.argmax(tf.reduce_max(outer, axis=2), axis=1)
self.yp2 = tf.argmax(tf.reduce_max(outer, axis=1), axis=1)
class Inference(object):
def __init__(self):
with open(word_emb_file, "r") as fh:
self.word_mat = np.array(json.load(fh), dtype=np.float32)
with open(char_emb_file, "r") as fh:
self.char_mat = np.array(json.load(fh), dtype=np.float32)
with open(word2idx_file, "r") as fh:
self.word2idx_dict = json.load(fh)
with open(char2idx_file, "r") as fh:
self.char2idx_dict = json.load(fh)
self.model = InfModel(self.word_mat, self.char_mat)
sess_config = tf.ConfigProto(allow_soft_placement=True)
sess_config.gpu_options.allow_growth = True
self.sess = tf.Session(config=sess_config)
saver = tf.train.Saver()
saver.restore(self.sess, tf.train.latest_checkpoint(save_dir))
def response(self, context, question):
sess = self.sess
model = self.model
span, context_idxs, ques_idxs, context_char_idxs, ques_char_idxs = \
self.prepro(context, question)
yp1, yp2 = \
sess.run(
[model.yp1, model.yp2],
feed_dict={
model.c: context_idxs, model.q: ques_idxs,
model.ch: context_char_idxs, model.qh: ques_char_idxs,
model.tokens_in_context: len(span)})
start_idx = span[yp1[0]][0]
end_idx = span[yp2[0]][1]
return context[start_idx: end_idx]
def prepro(self, context, question):
context = context.replace("''", '" ').replace("``", '" ')
context_tokens = word_tokenize(context)
context_chars = [list(token) for token in context_tokens]
spans = convert_idx(context, context_tokens)
ques = question.replace("''", '" ').replace("``", '" ')
ques_tokens = word_tokenize(ques)
ques_chars = [list(token) for token in ques_tokens]
context_idxs = np.zeros([1, len(context_tokens)], dtype=np.int32)
context_char_idxs = np.zeros(
[1, len(context_tokens), char_limit], dtype=np.int32)
ques_idxs = np.zeros([1, len(ques_tokens)], dtype=np.int32)
ques_char_idxs = np.zeros(
[1, len(ques_tokens), char_limit], dtype=np.int32)
def _get_word(word):
for each in (word, word.lower(), word.capitalize(), word.upper()):
if each in self.word2idx_dict:
return self.word2idx_dict[each]
return 1
def _get_char(char):
if char in self.char2idx_dict:
return self.char2idx_dict[char]
return 1
for i, token in enumerate(context_tokens):
context_idxs[0, i] = _get_word(token)
for i, token in enumerate(ques_tokens):
ques_idxs[0, i] = _get_word(token)
for i, token in enumerate(context_chars):
for j, char in enumerate(token):
if j == char_limit:
break
context_char_idxs[0, i, j] = _get_char(char)
for i, token in enumerate(ques_chars):
for j, char in enumerate(token):
if j == char_limit:
break
ques_char_idxs[0, i, j] = _get_char(char)
return spans, context_idxs, ques_idxs, context_char_idxs, ques_char_idxs
# Demo, example from paper "SQuAD: 100,000+ Questions for Machine Comprehension of Text"
if __name__ == "__main__":
infer = Inference()
context = "In meteorology, precipitation is any product of the condensation " \
"of atmospheric water vapor that falls under gravity. The main forms " \
"of precipitation include drizzle, rain, sleet, snow, graupel and hail." \
"Precipitation forms as smaller droplets coalesce via collision with other " \
"rain drops or ice crystals within a cloud. Short, intense periods of rain " \
"in scattered locations are called “showers”."
ques1 = "What causes precipitation to fall?"
ques2 = "What is another main form of precipitation besides drizzle, rain, snow, sleet and hail?"
ques3 = "Where do water droplets collide with ice crystals to form precipitation?"
# Correct: gravity, Output: drizzle, rain, sleet, snow, graupel and hail
ans1 = infer.response(context, ques1)
print("Answer 1: {}".format(ans1))
# Correct: graupel, Output: graupel
ans2 = infer.response(context, ques2)
print("Answer 2: {}".format(ans2))
# Correct: within a cloud, Output: within a cloud
ans3 = infer.response(context, ques3)
print("Answer 3: {}".format(ans3))