-
Notifications
You must be signed in to change notification settings - Fork 0
/
predict.py
79 lines (53 loc) · 2.1 KB
/
predict.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
import tensorflow as tf
import numpy as np
import os
import sys
import model
import data_process as data
from configs import DEFINES
def main(self):
arg_length = len(sys.argv)
if(arg_length < 2):
raise Exception("You should put one sentences to predict")
inputs = []
for i in sys.argv[1:]:
inputs.append(i)
inputs = " ".join(inputs)
_, _, t2i, i2t, max_len = data.load_data(DEFINES.data_path)
encoder_inputs, decoder_inputs = prepare_pred_input(inputs, t2i, max_len)
print(encoder_inputs, decoder_inputs)
embedding_matrix = data.get_embedding_matrix(DEFINES.data_path, DEFINES.embedding_path, i2t)
params = make_params(embedding_matrix, max_len)
estimator = tf.estimator.Estimator(model_fn = model.model_fn,
model_dir = DEFINES.check_point,
params = params)
predict_input_fn = tf.estimator.inputs.numpy_input_fn(
x={"encoer_inputs": encoder_inputs, "decoder_inputs": decoder_inputs},
num_epochs=1,
shuffle=False)
predict = estimator.predict(input_fn = predict_input_fn)
prediction = next(predict)['prediction']
print(inputs)
print(data.token2str(prediction, i2t))
def prepare_pred_input(inputs, t2i, max_len):
result = []
for token in inputs.split():
if token in t2i:
result.append(t2i[token])
else:
result.append(t2i['<UNK>'])
encoder_inputs = np.array([result+[0]*(max_len-len(result))])
decoder_inputs = np.array([[1]+[0]*(max_len)])
return encoder_inputs, decoder_inputs
def make_params(embedding_matrix, max_len):
params = {'hidden_dim': DEFINES.hidden_dim,
'embedding_matrix': embedding_matrix,
'vocab_size': embedding_matrix.shape[0],
'max_length': max_len,
'teacher_forcing_rate': DEFINES.teacher_forcing_rate,
'dropout_rate': DEFINES.dropout_rate,
'learning_rate': DEFINES.learning_rate}
return params
if __name__ == '__main__':
tf.logging.set_verbosity(tf.logging.INFO)
tf.app.run(main)