from __future__ import print_function

import tensorflow as tf
import time
import os
import pickle

from tf_lstm_char_seq_utils import TextLoader
from tf_lstm_char_seq_model import Model


class arguments:
    def __init__(self):
        return


def main():
    args = arguments()
    train(args)


def train(args):
    args.data_dir = 'char_seq'
    args.save_dir = 'char_seq'
    args.rnn_size = 64
    args.num_layers = 1
    args.batch_size = 50
    args.seq_length = 50
    args.num_epochs = 5
    args.save_every = 1000
    args.grad_clip = 5.
    args.learning_rate = 0.002
    args.decay_rate = 0.97

    data_loader = TextLoader(args.data_dir, args.batch_size, args.seq_length)
    args.vocab_size = data_loader.vocab_size

    with open(os.path.join(args.data_dir, 'config.pkl'), 'wb') as f:
        pickle.dump(args, f)

    with open(os.path.join(args.save_dir, 'chars_vocab.pkl'), 'wb') as f:
        pickle.dump((data_loader.chars, data_loader.vocab), f)

    model = Model(args)

    with tf.Session() as sess:
        tf.global_variables_initializer().run()
        saver = tf.train.Saver(tf.all_variables())
        for e in range(args.num_epochs):
            sess.run(tf.assign(model.lr, args.learning_rate * (args.decay_rate ** e)))
            data_loader.reset_batch_pointer()
            state = sess.run(model.initial_state)

            for b in range(data_loader.num_batches):
                start = time.time()
                x, y = data_loader.next_batch()
                feed = {model.input_data: x, model.targets: y}
                for i, (c, h) in enumerate(model.initial_state):
                    feed[c] = state[i].c
                    feed[h] = state[i].h

                train_loss, state, _ = sess.run([model.cost, model.final_state, model.train_op], feed)
                end = time.time()
                print("{}/{} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}".format(
                    e * data_loader.num_batches + b, args.num_epochs * data_loader.num_batches,
                    e, train_loss, end - start))

                if e == args.num_epochs - 1 and b == data_loader.num_batches - 1:
                    checkpoint_path = os.path.join(args.save_dir, 'model.ckpt')
                    saver.save(sess, checkpoint_path, global_step=e * data_loader.num_batches + b)
                    print("model saved to {}".format(checkpoint_path))


if __name__ == '__main__':
    main()
from __future__ import print_function

import tensorflow as tf
import time
import os
import pickle

from tf_lstm_char_seq_utils import TextLoader
from tf_lstm_char_seq_model import Model


class arguments:
    def __init__(self):
        return


def main():
    args = arguments()
    train(args)


def train(args):
    args.data_dir = 'char_seq'
    args.save_dir = 'char_seq'
    args.rnn_size = 64
    args.num_layers = 1
    args.batch_size = 50
    args.seq_length = 50
    args.num_epochs = 5
    args.save_every = 1000
    args.grad_clip = 5.
    args.learning_rate = 0.002
    args.decay_rate = 0.97

    data_loader = TextLoader(args.data_dir, args.batch_size, args.seq_length)
    args.vocab_size = data_loader.vocab_size

    with open(os.path.join(args.data_dir, 'config.pkl'), 'wb') as f:
        pickle.dump(args, f)

    with open(os.path.join(args.save_dir, 'chars_vocab.pkl'), 'wb') as f:
        pickle.dump((data_loader.chars, data_loader.vocab), f)

    model = Model(args)

    with tf.Session() as sess:
        tf.global_variables_initializer().run()
        saver = tf.train.Saver(tf.all_variables())
        for e in range(args.num_epochs):
            sess.run(tf.assign(model.lr, args.learning_rate * (args.decay_rate ** e)))
            data_loader.reset_batch_pointer()
            state = sess.run(model.initial_state)

            for b in range(data_loader.num_batches):
                start = time.time()
                x, y = data_loader.next_batch()
                feed = {model.input_data: x, model.targets: y}
                for i, (c, h) in enumerate(model.initial_state):
                    feed[c] = state[i].c
                    feed[h] = state[i].h

                train_loss, state, _ = sess.run([model.cost, model.final_state, model.train_op], feed)
                end = time.time()
                print("{}/{} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}".format(
                    e * data_loader.num_batches + b, args.num_epochs * data_loader.num_batches,
                    e, train_loss, end - start))

                if e == args.num_epochs - 1 and b == data_loader.num_batches - 1:
                    checkpoint_path = os.path.join(args.save_dir, 'model.ckpt')
                    saver.save(sess, checkpoint_path, global_step=e * data_loader.num_batches + b)
                    print("model saved to {}".format(checkpoint_path))


if __name__ == '__main__':
    main()