import tensorflow as tf
import numpy as np
tf.reset_default_graph()
sentences = ['i love damao','i like mengjun','we love all']
words = list(set(" ".join(sentences).split()))
word2idx = {v:k for k,v in enumerate(words)}
idx2word = {k:v for k,v in enumerate(words)}
V = len(words)   # 詞典大小
step = 2   # 時間序列長度
hidden = 5   # 隱層大小
dim = 50   # 詞向量維度
# 制作輸入和標簽
def make_batch(sentences):
    input_batch = []
    target_batch = []
    for sentence in sentences:
        words = sentence.split()
        input = [word2idx[word] for word in words[:-1]]
        target = word2idx[words[-1]]
        input_batch.append(input)
        target_batch.append(np.eye(V)[target])   # 這里將標簽改為 one-hot 編碼，之后計算交叉熵的時候會用到
    return input_batch, target_batch
# 初始化詞向量
embedding = tf.get_variable(shape=[V, dim], initializer=tf.random_normal_initializer(), name="embedding")
X = tf.placeholder(tf.int32, [None, step])
XX = tf.nn.embedding_lookup(embedding,  X)
Y = tf.placeholder(tf.int32, [None, V])
# 定義 cell
cell = tf.nn.rnn_cell.BasicRNNCell(hidden)
# 計算各個時間點的輸出和隱層輸出的結(jié)果
outputs, hiddens = tf.nn.dynamic_rnn(cell, XX, dtype=tf.float32)     # outputs: [batch_size, step, hidden] hiddens: [batch_size, hidden]
# 這里將所有時間點的狀態(tài)向量都作為了后續(xù)分類器的輸入（也可以只將最后時間節(jié)點的狀態(tài)向量作為后續(xù)分類器的輸入）
W = tf.Variable(tf.random_normal([step*hidden, V]))
b = tf.Variable(tf.random_normal([V]))
L = tf.matmul(tf.reshape(outputs,[-1, step*hidden]), W) + b
# 計算損失并進行優(yōu)化
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(labels=Y, logits=L))
optimizer = tf.train.AdamOptimizer(0.001).minimize(cost)
# 預(yù)測
prediction = tf.argmax(L, 1)
# 初始化 tf
init = tf.global_variables_initializer()
sess = tf.Session()
sess.run(init)
# 喂訓練數(shù)據(jù)
input_batch, target_batch = make_batch(sentences)
for epoch in range(5000):
    _, loss = sess.run([optimizer, cost], feed_dict={X:input_batch, Y:target_batch})
    if (epoch+1)%1000 == 0:
        print("epoch: ", '%04d'%(epoch+1), 'cost= ', '%04f'%(loss))
# 預(yù)測數(shù)據(jù)
predict = sess.run([prediction], feed_dict={X: input_batch})
print([sentence.split()[:2] for sentence in sentences], '->', [idx2word[n] for n in predict[0]])