from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import numpy as np
import tensorflow as tf
import random
import argparse
import math
from colors import *
from tqdm import *
from handler import Batch, DatasetBase, DatasetTrain, DatasetEval, DatasetTest
import util
from gensim.models import Word2Vec
FLAGS = None
random.seed(0)
np.random.seed(0)
tf.set_random_seed(0)
#filename = '/tmp'
#total_line_num = 100
#train_line_num = 80
#eval_line_num = 20
#filename = '/xaa'
#total_line_num = 50000
#train_line_num = 45000
#eval_line_num = 5000
#filename = '/clr_conversation.txt'
#total_line_num = 2842478
#train_line_num = 2840000
#eval_line_num = 2478
filename = '/conversation.txt'
total_line_num = 3599478
train_line_num = 3587000
eval_line_num = 12478
emb_size = 300
PKL_EXIST = True
MAX_SENTENCE_LENGTH = 15 # longest
special_tokens = {'<PAD>': 0, '<BOS>': 1, '<EOS>': 2, '<UNK>': 3}
special_tokens_to_word = ['<PAD>', '<BOS>', '<EOS>', '<UNK>']
modes = {'train': 0, 'eval': 1, 'test': 2}
class Seq2Seq:
def __init__(self, voc, idx2word, mode, att, lr=None):
self.num_layers = 2
self.rnn_size = 2048
self.keep_prob = 1.0
self.vocab_num = voc
self.with_attention = att
self.mode = mode
self.lr = lr
self.idx2word = idx2word
self.embedding_size = emb_size
self.hist_summary = []
def _create_rnn_cell(self):
def single_rnn_cell():
cell = tf.contrib.rnn.GRUCell(self.rnn_size)
#cell = tf.contrib.rnn.LSTMCell(self.rnn_size, initializer=tf.orthogonal_initializer())
if self.mode == modes['train']:
cell = tf.contrib.rnn.DropoutWrapper(cell, input_keep_prob=self.keep_prob,
output_keep_prob=self.keep_prob)
return cell
cell = tf.contrib.rnn.MultiRNNCell([single_rnn_cell() for _ in range(self.num_layers)])
return cell
def build_model(self, emb=None):
self.encoder_inputs = tf.placeholder(tf.int32, [None, None], name='encoder_inputs')
self.encoder_inputs_length = tf.placeholder(tf.int32, [None], name='encoder_inputs_length')
with tf.device("/cpu:0"):
embed = tf.constant_initializer(emb, dtype=tf.float32)
embedding = tf.get_variable(
initializer=embed, shape=emb.shape, dtype=tf.float32, trainable=True, name='embedding')
#embedding = tf.get_variable('embedding', [self.vocab_num, self.embedding_size])
self.hist_summary.append(tf.summary.histogram(embedding.name + '/embed', embedding))
self.batch_size = tf.placeholder(tf.int32, [], name='batch_size')
self.sampling_prob = tf.placeholder(tf.float32, [], name='sampling_prob')
if self.mode != modes['test']:
self.decoder_targets = tf.placeholder(tf.int32, [None, None], name='decoder_targets')
self.decoder_targets_length = tf.placeholder(tf.int32, [None], name='decoder_targets_length')
self.max_target_sequence_length = tf.reduce_max(self.decoder_targets_length, name='max_target_len')
self.mask = tf.sequence_mask(self.decoder_targets_length, self.max_target_sequence_length,
dtype=tf.float32, name='masks')
with tf.variable_scope('encoder'):
encoder_cell = self._create_rnn_cell()
with tf.device("/cpu:0"):
encoder_inputs_embedded = tf.nn.embedding_lookup(embedding, self.encoder_inputs)
encoder_outputs, encoder_state = tf.nn.dynamic_rnn(encoder_cell, encoder_inputs_embedded,
sequence_length=self.encoder_inputs_length,
dtype=tf.float32)
with tf.variable_scope('decoder'):
encoder_inputs_length = self.encoder_inputs_length
decoder_cell = self._create_rnn_cell()
batch_size = self.batch_size
if self.with_attention:
print('wrapped with bahdanau attention...')
attention_mechanism = tf.contrib.seq2seq.BahdanauAttention(
num_units=self.rnn_size, memory=encoder_outputs,
memory_sequence_length=encoder_inputs_length)
decoder_cell = tf.contrib.seq2seq.AttentionWrapper(
cell=decoder_cell, attention_mechanism=attention_mechanism,
attention_layer_size=self.rnn_size, name='Attention_Wrapper')
decoder_initial_state = decoder_cell.zero_state(batch_size=batch_size,
dtype=tf.float32).clone(cell_state=encoder_state)
else:
decoder_initial_state = encoder_state
projection_layer = tf.layers.Dense(
self.vocab_num, kernel_initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1))
#self.vocab_num, kernel_initializer=tf.contrib.layers.xavier_initializer())
if self.mode == modes['train']:
ending = tf.strided_slice(self.decoder_targets, [0, 0], [self.batch_size, -1], [1, 1])
decoder_input = tf.concat([tf.fill([self.batch_size, 1], special_tokens['<BOS>']), ending], 1)
with tf.device("/cpu:0"):
decoder_inputs_embedded = tf.nn.embedding_lookup(embedding, decoder_input)
training_helper = tf.contrib.seq2seq.ScheduledEmbeddingTrainingHelper(
inputs=decoder_inputs_embedded,
sequence_length=self.decoder_targets_length,
embedding=embedding,
time_major=False,
sampling_probability=self.sampling_prob,
name='training_helper')
training_decoder = tf.contrib.seq2seq.BasicDecoder(cell=decoder_cell, helper=training_helper,
initial_state=decoder_initial_state, output_layer=projection_layer)
decoder_outputs, _, final_seq_len = tf.contrib.seq2seq.dynamic_decode(decoder=training_decoder,
impute_finished=True,
maximum_iterations=self.max_target_sequence_length)
#pad_size = self.max_target_sequence_length - tf.reduce_max(final_seq_len)
#pad_rnn_output = tf.pad(decoder_outputs.rnn_output, [[0, 0], [0, pad_size], [0, 0]])
#self.decoder_logits_train = tf.identity(pad_rnn_output)
self.decoder_logits_train = tf.identity(decoder_outputs.rnn_output)
self.decoder_predict_train = tf.argmax(self.decoder_logits_train, axis=-1, name='decoder_pred_train')
print(self.decoder_predict_train)
self.train_loss = tf.contrib.seq2seq.sequence_loss(logits=self.decoder_logits_train,
targets=self.decoder_targets, weights=self.mask)
self.train_summary = tf.summary.scalar('training loss', self.train_loss)
elif self.mode == modes['eval']:
start_tokens = tf.ones([self.batch_size, ], tf.int32) # * special_tokens['<BOS>']
end_token = special_tokens['<EOS>']
decoding_helper = tf.contrib.seq2seq.GreedyEmbeddingHelper(embedding=embedding,
start_tokens=start_tokens, end_token=end_token)
inference_decoder = tf.contrib.seq2seq.BasicDecoder(cell=decoder_cell, helper=decoding_helper,
initial_state=decoder_initial_state,
output_layer=projection_layer)
decoder_outputs, _, final_seq_len = tf.contrib.seq2seq.dynamic_decode(decoder=inference_decoder,
maximum_iterations=self.max_target_sequence_length)
# pad to same shape in order to calculate loss
pad_size = self.max_target_sequence_length - tf.reduce_max(final_seq_len)
pad_rnn_output = tf.pad(decoder_outputs.rnn_output, [[0, 0], [0, pad_size], [0, 0]])
# self.decoder_logits_eval = tf.identity(decoder_outputs.rnn_output)
self.decoder_logits_eval = tf.identity(pad_rnn_output)
self.decoder_predict_eval = tf.argmax(self.decoder_logits_eval, axis=-1, name='decoder_pred_eval')
print(self.decoder_predict_eval)
self.eval_loss = tf.contrib.seq2seq.sequence_loss(logits=pad_rnn_output,
targets=self.decoder_targets, weights=self.mask)
self.eval_summary = tf.summary.scalar('validation loss', self.eval_loss)
elif self.mode == modes['test']:
start_tokens = tf.ones([self.batch_size, ], tf.int32) # * special_tokens['<BOS>']
end_token = special_tokens['<EOS>']
decoding_helper = tf.contrib.seq2seq.GreedyEmbeddingHelper(embedding=embedding,
start_tokens=start_tokens, end_token=end_token)
inference_decoder = tf.contrib.seq2seq.BasicDecoder(cell=decoder_cell, helper=decoding_helper,
initial_state=decoder_initial_state,
output_layer=projection_layer)
decoder_outputs, _, final_seq_len = tf.contrib.seq2seq.dynamic_decode(decoder=inference_decoder,
maximum_iterations=MAX_SENTENCE_LENGTH)
self.decoder_logits_eval = tf.identity(decoder_outputs.rnn_output)
self.decoder_predict_eval = tf.argmax(self.decoder_logits_eval, axis=-1, name='decoder_pred_eval')
print(self.decoder_predict_eval)
def build_optimizer(self, lr):
optimizer = tf.train.GradientDescentOptimizer(lr)
#optimizer = tf.train.AdamOptimizer(lr)#.minimize(self.train_loss)
#optimizer = tf.train.AdamOptimizer(lr, epsilon=0.1)#.minimize(self.train_loss)
trainable_params = tf.trainable_variables()
gradients = tf.gradients(self.train_loss, trainable_params)
gradients, _ = tf.clip_by_global_norm(gradients, 5.0)
gradients = list(zip(gradients, trainable_params))
self.train_op = optimizer.apply_gradients(gradients)
for grad, var in gradients:
self.hist_summary.append(tf.summary.histogram(var.name + '/gradient', grad))
self.hist_summary.append(self.train_summary)
self.train_summary = tf.summary.merge(self.hist_summary)
def train(self, sess, batch, print_pred, summary_writer, add_global, prob):
feed_dict = {self.encoder_inputs: batch.encoder_inputs,
self.encoder_inputs_length: batch.encoder_inputs_length,
self.decoder_targets: batch.decoder_targets,
self.decoder_targets_length: batch.decoder_targets_length,
self.batch_size: len(batch.encoder_inputs),
self.sampling_prob: prob}
if print_pred:
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
_, loss, pred, summary, current_step, print_lr = sess.run([self.train_op, self.train_loss,
self.decoder_predict_train, self.train_summary, add_global, self.lr],
feed_dict=feed_dict, options=run_options)
i = np.random.randint(0, len(batch.encoder_inputs))
util.decoder_print(self.idx2word, batch.encoder_inputs[i], batch.encoder_inputs_length[i],
batch.decoder_targets[i], batch.decoder_targets_length[i], pred[i], 'yellow')
summary_writer.add_summary(summary, global_step=current_step)
else:
_, loss, current_step, print_lr = sess.run([self.train_op, self.train_loss,
add_global, self.lr], feed_dict=feed_dict)
return loss, calc_perplexity(loss), current_step, print_lr
def eval(self, sess, batch, summary_writer, current_step):
feed_dict = {self.encoder_inputs: batch.encoder_inputs,
self.encoder_inputs_length: batch.encoder_inputs_length,
self.decoder_targets: batch.decoder_targets,
self.decoder_targets_length: batch.decoder_targets_length,
self.batch_size: len(batch.encoder_inputs)}
loss, pred, summary = sess.run([self.eval_loss,
self.decoder_predict_eval, self.eval_summary], feed_dict=feed_dict)
print_num = 3
print_more = random.sample(range(len(batch.encoder_inputs)), print_num)
#repeat bug: np.random.randint(len(batch.encoder_inputs), size=(print_num))
for i in print_more:
util.decoder_print(self.idx2word, batch.encoder_inputs[i], batch.encoder_inputs_length[i],
batch.decoder_targets[i], batch.decoder_targets_length[i], pred[i], 'green')
summary_writer.add_summary(summary, global_step=current_step)
return loss, calc_perplexity(loss)
def inference(self, sess, batch, txt):
feed_dict = {self.encoder_inputs: batch.encoder_inputs,
self.encoder_inputs_length: batch.encoder_inputs_length,
self.batch_size: len(batch.encoder_inputs)}
pred = sess.run([self.decoder_predict_eval], feed_dict=feed_dict)
pred = pred[0]
for i in range(len(pred)):
sent = util.decoder_inference(self.idx2word, batch.encoder_inputs[i],
batch.encoder_inputs_length[i], pred[i])
txt.write(sent + "\n")
def calc_perplexity(loss):
return math.exp(float(loss)) if loss < 300 else float('inf')
def train():
datasetTrain = DatasetTrain()
print('start build dict...')
train_data, eval_data = datasetTrain.build_dict(FLAGS.data_dir, filename,
FLAGS.min_counts, train_line_num, eval_line_num, emb_size, PKL_EXIST)
print('build dict done!')
datasetTrain.prep(train_data)
datasetEval = DatasetEval()
datasetEval.load_dict()
datasetEval.prep(eval_data)
word2vec_model = datasetTrain.model
embeddings = np.zeros([datasetTrain.vocab_num, emb_size],dtype=np.float32)
for word in word2vec_model.wv.vocab:
index = word2vec_model.wv.vocab[word].index
vec = word2vec_model.wv[word]
embeddings[index] = vec
if word in special_tokens_to_word:
print('special word: ', word, ', index: ', index)
pad = np.random.normal(size=[emb_size])
unk = np.random.normal(size=[emb_size])
embeddings[0] = pad
embeddings[3] = unk
#print(embeddings)
# np.save('embeddings.npy', embeddings)
train_graph = tf.Graph()
eval_graph = tf.Graph()
gpu_config = tf.ConfigProto()
gpu_config.gpu_options.allow_growth = True
print('start building train graph...')
with train_graph.as_default():
global_step = tf.Variable(0, trainable=False)
lr = tf.train.exponential_decay(FLAGS.learning_rate,
global_step=global_step,
decay_steps=3250, decay_rate=0.95, staircase=True)
add_global = global_step.assign_add(1)
model = Seq2Seq(voc=datasetTrain.vocab_num, idx2word=datasetTrain.idx2word,
mode=modes['train'], att=FLAGS.with_attention, lr=lr)
model.build_model(embeddings)
model.build_optimizer(FLAGS.learning_rate)
model.saver = tf.train.Saver(max_to_keep = 3)
init = tf.global_variables_initializer()
train_sess = tf.Session(graph=train_graph, config=gpu_config)
print('start building eval graph...')
with eval_graph.as_default():
model_eval = Seq2Seq(voc=datasetEval.vocab_num, idx2word=datasetEval.idx2word,
mode=modes['eval'], att=FLAGS.with_attention)
model_eval.build_model(embeddings)
model_eval.saver = tf.train.Saver(max_to_keep = 3)
#init_eval = tf.global_variables_initializer()
eval_sess = tf.Session(graph=eval_graph, config=gpu_config)
ckpt = tf.train.get_checkpoint_state(FLAGS.save_dir)
if FLAGS.load_saver and ckpt and tf.train.checkpoint_exists(ckpt.model_checkpoint_path):
print('Reloading model parameters..')
model.saver.restore(train_sess, ckpt.model_checkpoint_path)
#model_eval.saver.restore(eval_sess, ckpt.model_checkpoint_path)
print(ckpt.model_checkpoint_path)
else:
print('Created new model parameters..')
train_sess.run(init)
#eval_sess.run(init_eval)
ckpts_path = FLAGS.save_dir + "chatbot.ckpt"
summary_writer = tf.summary.FileWriter(FLAGS.log_dir + '/train')
summary_writer.add_graph(train_graph)
summary_writer.add_graph(eval_graph)
num_steps = int( len(datasetTrain.data) / FLAGS.batch_size )
pbar = tqdm(range(FLAGS.num_epochs))
pt = 0
total_samp = FLAGS.num_epochs * 3
samp_prob = util.inv_sigmoid(total_samp)
current_step = 0 # after first round, assign global_step to current_step
for epo in pbar:
print(color('start: epoch ' + str(epo), fg='blue', bg='white'))
for i in range(num_steps):
batch = datasetTrain.next_batch(FLAGS.batch_size, shuffle=True)
print_pred = False
if current_step % FLAGS.num_display_steps == 0 and current_step != 0:
print_pred = True
loss, perp, current_step, print_lr = model.train(train_sess, batch, print_pred,
summary_writer, add_global, samp_prob[pt])
if current_step % FLAGS.num_saver_steps == 0 and current_step != 0:
ckpt_path = model.saver.save(train_sess, ckpts_path, global_step=current_step)
print(color("\nSaver saved: " + ckpt_path, fg='white', bg='green', style='bold'))
model_eval.saver.restore(eval_sess, ckpt_path)
print(color("\n[Eval. Prediction] Epoch " + str(epo) + ", step " + str(i) + "/" \
+ str(num_steps) + "......", fg='white', bg='green', style='underline'))
batch_eval = datasetEval.next_batch(FLAGS.batch_size, shuffle=True)
loss_eval, perp_eval = model_eval.eval(eval_sess, batch_eval, summary_writer, current_step)
print(color("Epoch " + str(epo) + ", step " + str(i) + "/" + str(num_steps) + \
", (Evaluation Loss: " + "{:.4f}".format(loss_eval) + \
", Perplexity: " + "{:.4f}".format(perp_eval) + ")", fg='white', bg='green'))
pbar.set_description("S " + str(i) + "/" + \
str(num_steps) + "(" + str(current_step) + ")" + \
", (L: " + "{:.4f}".format(loss) + ", P: " + "{:.1f}".format(perp) + ", Samp: "+ \
"{:.4f}".format(samp_prob[pt]) + ", lr: "+ "{:.6f}".format(print_lr) + ")" )
if i % int(num_steps / 3) == 0 and i != 0:
pt += 1
print(color('sampling pt: ' + str( pt ) + '/' + str(total_samp), fg='white', bg='red'))
print(color('sampling pt: ' + str( pt ) + '/' + str(total_samp), fg='white', bg='red'))
def test():
datasetTest = DatasetTest()
print('start build dict...')
datasetTest.load_dict()
print('build dict done!')
test_data = datasetTest.load_test_line(FLAGS.test_dir, FLAGS.test_file)
datasetTest.prep(test_data)
# word2vec_model = datasetTest.model
embeddings = np.zeros([datasetTest.vocab_num, emb_size],dtype=np.float32)
# for word in word2vec_model.wv.vocab:
# index = word2vec_model.wv.vocab[word].index
# vec = word2vec_model.wv[word]
# embeddings[index] = vec
# if word in special_tokens_to_word:
# print('special word: ', word, ', index: ', index)
# pad = np.random.normal(size=[emb_size])
# unk = np.random.normal(size=[emb_size])
# embeddings[0] = pad
# embeddings[3] = unk
#print(embeddings)
test_graph = tf.Graph()
gpu_config = tf.ConfigProto()
gpu_config.gpu_options.allow_growth = True
print('start building test graph...')
with test_graph.as_default():
model_test = Seq2Seq(voc=datasetTest.vocab_num, idx2word=datasetTest.idx2word,
mode=modes['test'], att=FLAGS.with_attention)
model_test.build_model(embeddings)
model_test.saver = tf.train.Saver(max_to_keep = 3)
test_sess = tf.Session(graph=test_graph, config=gpu_config)
ckpt = tf.train.get_checkpoint_state(FLAGS.save_dir)
if FLAGS.load_saver and ckpt and tf.train.checkpoint_exists(ckpt.model_checkpoint_path):
print('Reloading model parameters..')
model_test.saver.restore(test_sess, ckpt.model_checkpoint_path)
print(ckpt.model_checkpoint_path)
else:
print('ERROR: you should load model for testing!')
exit(0)
ckpts_path = FLAGS.save_dir + "chatbot.ckpt"
# You have to run this write_graph to get your testing model!
# (since that train/test model structure are different)
tf.train.write_graph(test_sess.graph_def, "./load", "test.pb", False) #proto
#exit(0)
num_steps = int( len(datasetTest.test_data) / FLAGS.batch_size )
txt = open(FLAGS.output_filename, 'w')
for i in range(num_steps):
batch = datasetTest.next_batch(FLAGS.batch_size)
model_test.inference(test_sess, batch, txt)
print(str(i) + '/' + str(num_steps) + ' steps...done')
print('\n\nTesting finished.')
print('\nSave file: ' + FLAGS.output_filename)
txt.close()
def main(_):
if FLAGS.test_mode == False:
#print(color('remove directory: ' + FLAGS.log_dir, fg='red'))
#if tf.gfile.Exists(FLAGS.log_dir):
# tf.gfile.DeleteRecursively(FLAGS.log_dir)
#tf.gfile.MakeDirs(FLAGS.log_dir)
print('train mode: start')
train()
else:
if FLAGS.load_saver == True:
print('load saver!!')
else:
print('ERROR: you cannot run test without saver...')
exit(0)
print('test mode: start')
test()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('-lr', '--learning_rate', type=float, default=0.005)
parser.add_argument('-mi', '--min_counts', type=int, default=100) #50 -> 15000 words
parser.add_argument('-e', '--num_epochs', type=int, default=30)
parser.add_argument('-b', '--batch_size', type=int, default=250)
parser.add_argument('-t', '--test_mode', type=int, default=0)
parser.add_argument('-d', '--num_display_steps', type=int, default=120)
parser.add_argument('-ns', '--num_saver_steps', type=int, default=200)
parser.add_argument('-s', '--save_dir', type=str, default='save/')
parser.add_argument('-l', '--log_dir', type=str, default='logs/')
parser.add_argument('-o', '--output_filename', type=str, default='output.txt')
parser.add_argument('-lo', '--load_saver', type=int, default=0)
parser.add_argument('-at', '--with_attention', type=int, default=1)
parser.add_argument('--data_dir', type=str,
default=('/home/data/mlds_hw2_2_data')
)
parser.add_argument('--test_dir', type=str,
default=('/home/data/mlds_hw2_2_data')
)
parser.add_argument('--test_file', type=str,
default=('/test_input.txt')
)
FLAGS, unparsed = parser.parse_known_args()
tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
