RNN_ChattingRobot_Week/test.py at master · AIBigTruth/RNN_ChattingRobot_Week · GitHub

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
import tensorflow as tf
import numpy as np
import os
from six.moves import xrange
_buckets = []
convo_hist_limit = 1
max_source_length = 0
max_target_length = 0
flags = tf.app.flags
FLAGS = flags.FLAGS

import datautil
import seq2seq_model


tf.reset_default_graph()
max_train_data_size= 0#(0: no limit)


dropout = 1.0
grad_clip = 5.0
batch_size = 60
hidden_size = 14
num_layers =2
learning_rate =0.5
lr_decay_factor =0.99


hidden_size = 100
checkpoint_dir= "datacn/checkpoints/"
_buckets =[(5, 5), (10, 10), (20, 20)]
def getdialogInfo():
    vocabch, rev_vocabch=datautil.initialize_vocabulary(os.path.join(datautil.data_dir, datautil.vocabulary_filech))
    vocab_sizech= len(vocabch)
    print("vocab_sizech",vocab_sizech)
    filesfrom,_=datautil.getRawFileList(datautil.data_dir+"fromids/")
    filesto,_=datautil.getRawFileList(datautil.data_dir+"toids/")
    source_train_file_path = filesfrom[0]
    target_train_file_path= filesto[0]
    return vocab_sizech,vocab_sizech,vocabch,rev_vocabch


def main():

    vocab_sizeen,vocab_sizech,vocaben,rev_vocabch= getdialogInfo()

    if not os.path.exists(checkpoint_dir):
        os.mkdir(checkpoint_dir)
    print ("checkpoint_dir is {0}".format(checkpoint_dir))

    with tf.Session() as sess:
        model = createModel(sess,True,vocab_sizeen,vocab_sizech)

        print (_buckets)
        model.batch_size = 1

        conversation_history =[]
        while True:
            prompt = "请输入: "
            sentence = input(prompt)
            conversation_history.append(sentence.strip())
            conversation_history = conversation_history[-convo_hist_limit:]

            token_ids = list(reversed( datautil.sentence_to_ids(" ".join(conversation_history) ,vocaben,normalize_digits=True,Isch=True) )  )
            #token_ids = list(reversed(vocab.tokens2Indices(" ".join(conversation_history))))
            print(token_ids)
            #token_ids = list(reversed(vocab.tokens2Indices(sentence)))
            bucket_id = min([b for b in xrange(len(_buckets))if _buckets[b][0] > len(token_ids)])

            encoder_inputs, decoder_inputs, target_weights = model.get_batch({bucket_id: [(token_ids, [])]}, bucket_id)

            _, _, output_logits = model.step(sess, encoder_inputs, decoder_inputs,target_weights, bucket_id, True)

            #TODO implement beam search
            outputs = [int(np.argmax(logit, axis=1)) for logit in output_logits]
            print("outputs",outputs,datautil.EOS_ID)
            if datautil.EOS_ID in outputs:
                outputs = outputs[:outputs.index(datautil.EOS_ID)]
                #print(vocab.indices2Tokens(outputs))
                #print("结果",datautil.ids2texts(outputs,rev_vocabch))
                convo_output =  " ".join(datautil.ids2texts(outputs,rev_vocabch))
                conversation_history.append(convo_output)
                print (convo_output)
            else:
                print("can not translation！")


def createModel(session, forward_only,from_vocab_size,to_vocab_size):
    """Create translation model and initialize or load parameters in session."""
    model = seq2seq_model.Seq2SeqModel(
      from_vocab_size,#from
      to_vocab_size,#to
      _buckets,
      hidden_size,
      num_layers,
      dropout,
      grad_clip,
      batch_size,
      learning_rate,
      lr_decay_factor,
      forward_only=forward_only,
      dtype=tf.float32)

    print("model is ok")


    ckpt = tf.train.latest_checkpoint(checkpoint_dir)
    if ckpt!=None:
        model.saver.restore(session, ckpt)
        print ("Reading model parameters from {0}".format(ckpt))
    else:
        print ("Created model with fresh parameters.")
        session.run(tf.global_variables_initializer())

    return model


if __name__=="__main__":
	main()