Merge branch 'release/v1.1.0'

.gitignore

@@ -0,0 +1,9 @@
+checkpoints/
+model/
+.idea/
+__pycache__/
+/.settings
+/.vs
+/data
+/.pydevproject
+/.project

README.md

@@ -0,0 +1,9 @@
+## Quick Start
+
+```
+git clone https://github.com/jinfagang/tensorflow_poems
+# train on poems
+python3 train.py
+# compose poems
+python3 compose_poem.py
+```

compose_poem.py

@@ -0,0 +1,75 @@
+import tensorflow as tf
+from poem.model import rnn_model
+from poem.poem import process_poem
+import numpy as np
+
+start_token = 'B'
+end_token = 'E'
+model_dir = './model/'
+corpus_file = './data/poem.txt'
+
+lr = 0.0002
+
+
+def to_word(predict, vocabs):
+    t = np.cumsum(predict)
+    s = np.sum(predict)
+    sample = int(np.searchsorted(t, np.random.rand(1) * s))
+    if sample > len(vocabs):
+        sample = len(vocabs) - 1
+    return vocabs[sample]
+
+
+def gen_poem(begin_word):
+    batch_size = 1
+    print('## loading corpus from %s' % model_dir)
+    poem_vector, word_int_map, vocabularies = process_poem(corpus_file)
+
+    input_data = tf.placeholder(tf.int32, [batch_size, None])
+
+    end_points = rnn_model(model='lstm', input_data=input_data, output_data=None, vocab_size=len(
+        vocabularies), rnn_size=128, num_layers=2, batch_size=64, learning_rate=lr)
+
+    saver = tf.train.Saver(tf.global_variables())
+    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
+    with tf.Session() as sess:
+        sess.run(init_op)
+
+        checkpoint = tf.train.latest_checkpoint(model_dir)
+        saver.restore(sess, checkpoint)
+
+        x = np.array([list(map(word_int_map.get, start_token))])
+
+        [predict, last_state] = sess.run([end_points['prediction'], end_points['last_state']],
+                                         feed_dict={input_data: x})
+        if begin_word:
+            word = begin_word
+        else:
+            word = to_word(predict, vocabularies)
+        poem_ = ''
+
+        i = 0
+        while word != end_token:
+            poem_ += word
+            i += 1
+            if i >= 24:
+                break
+            x = np.zeros((1, 1))
+            x[0, 0] = word_int_map[word]
+            [predict, last_state] = sess.run([end_points['prediction'], end_points['last_state']],
+                                             feed_dict={input_data: x, end_points['initial_state']: last_state})
+            word = to_word(predict, vocabularies)
+
+        return poem_
+
+
+def pretty_print_poem(poem_):
+    poem_sentences = poem_.split('。')
+    for s in poem_sentences:
+        if s != '' and len(s) > 10:
+            print(s + '。')
+
+if __name__ == '__main__':
+    begin_char = input('## please input the first character:')
+    poem = gen_poem(begin_char)
+    pretty_print_poem(poem_=poem)

data/数据介绍.txt

@@ -0,0 +1,3 @@
+## 下载
+
+https://pan.baidu.com/

src/compose_poem.py

@@ -0,0 +1,75 @@
+import tensorflow as tf
+from poem.model import rnn_model
+from poem.poem import process_poem
+import numpy as np
+
+start_token = 'B'
+end_token = 'E'
+model_dir = './model/'
+corpus_file = './data/poem.txt'
+
+lr = 0.0002
+
+
+def to_word(predict, vocabs):
+    t = np.cumsum(predict)
+    s = np.sum(predict)
+    sample = int(np.searchsorted(t, np.random.rand(1) * s))
+    if sample > len(vocabs):
+        sample = len(vocabs) - 1
+    return vocabs[sample]
+
+
+def gen_poem(begin_word):
+    batch_size = 1
+    print('## loading corpus from %s' % model_dir)
+    poem_vector, word_int_map, vocabularies = process_poem(corpus_file)
+
+    input_data = tf.placeholder(tf.int32, [batch_size, None])
+
+    end_points = rnn_model(model='lstm', input_data=input_data, output_data=None, vocab_size=len(
+        vocabularies), rnn_size=128, num_layers=2, batch_size=64, learning_rate=lr)
+
+    saver = tf.train.Saver(tf.global_variables())
+    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
+    with tf.Session() as sess:
+        sess.run(init_op)
+
+        checkpoint = tf.train.latest_checkpoint(model_dir)
+        saver.restore(sess, checkpoint)
+
+        x = np.array([list(map(word_int_map.get, start_token))])
+
+        [predict, last_state] = sess.run([end_points['prediction'], end_points['last_state']],
+                                         feed_dict={input_data: x})
+        if begin_word:
+            word = begin_word
+        else:
+            word = to_word(predict, vocabularies)
+        poem_ = ''
+
+        i = 0
+        while word != end_token:
+            poem_ += word
+            i += 1
+            if i >= 24:
+                break
+            x = np.zeros((1, 1))
+            x[0, 0] = word_int_map[word]
+            [predict, last_state] = sess.run([end_points['prediction'], end_points['last_state']],
+                                             feed_dict={input_data: x, end_points['initial_state']: last_state})
+            word = to_word(predict, vocabularies)
+
+        return poem_
+
+
+def pretty_print_poem(poem_):
+    poem_sentences = poem_.split('。')
+    for s in poem_sentences:
+        if s != '' and len(s) > 10:
+            print(s + '。')
+
+if __name__ == '__main__':
+    begin_char = input('## please input the first character:')
+    poem = gen_poem(begin_char)
+    pretty_print_poem(poem_=poem)

src/poem/model.py

@@ -0,0 +1,73 @@
+import tensorflow as tf
+import numpy as np
+
+
+def rnn_model(model, input_data, output_data, vocab_size, rnn_size=128, num_layers=2, batch_size=64,
+              learning_rate=0.01):
+    """
+    construct rnn seq2seq model.
+    :param model: model class
+    :param input_data: input data placeholder
+    :param output_data: output data placeholder
+    :param vocab_size:
+    :param rnn_size:
+    :param num_layers:
+    :param batch_size:
+    :param learning_rate:
+    :return:
+    """
+    end_points = {}
+
+    if model == 'rnn':
+        cell_fun = tf.contrib.rnn.BasicRNNCell
+    elif model == 'gru':
+        cell_fun = tf.contrib.rnn.GRUCell
+    elif model == 'lstm':
+        cell_fun = tf.contrib.rnn.BasicLSTMCell
+
+    cell = cell_fun(rnn_size, state_is_tuple=True)
+    cell = tf.contrib.rnn.MultiRNNCell([cell] * num_layers, state_is_tuple=True)
+
+    if output_data is not None:
+        initial_state = cell.zero_state(batch_size, tf.float32)
+    else:
+        initial_state = cell.zero_state(1, tf.float32)
+
+    with tf.device("/cpu:0"):
+        embedding = tf.get_variable('embedding', initializer=tf.random_uniform(
+            [vocab_size + 1, rnn_size], -1.0, 1.0))
+        inputs = tf.nn.embedding_lookup(embedding, input_data)
+
+    # [batch_size, ?, rnn_size] = [64, ?, 128]
+    outputs, last_state = tf.nn.dynamic_rnn(cell, inputs, initial_state=initial_state)
+    output = tf.reshape(outputs, [-1, rnn_size])
+
+    weights = tf.Variable(tf.truncated_normal([rnn_size, vocab_size + 1]))
+    bias = tf.Variable(tf.zeros(shape=[vocab_size + 1]))
+    logits = tf.nn.bias_add(tf.matmul(output, weights), bias=bias)
+    # [?, vocab_size+1]
+
+    if output_data is not None:
+        # output_data must be one-hot encode
+        labels = tf.one_hot(tf.reshape(output_data, [-1]), depth=vocab_size + 1)
+        # should be [?, vocab_size+1]
+
+        loss = tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=logits)
+        # loss shape should be [?, vocab_size+1]
+        total_loss = tf.reduce_mean(loss)
+        train_op = tf.train.AdamOptimizer(learning_rate).minimize(total_loss)
+
+        end_points['initial_state'] = initial_state
+        end_points['output'] = output
+        end_points['train_op'] = train_op
+        end_points['total_loss'] = total_loss
+        end_points['loss'] = loss
+        end_points['last_state'] = last_state
+    else:
+        prediction = tf.nn.softmax(logits)
+
+        end_points['initial_state'] = initial_state
+        end_points['last_state'] = last_state
+        end_points['prediction'] = prediction
+
+    return end_points

src/poem/poem.py

@@ -0,0 +1,65 @@
+import collections
+import os
+import sys
+import numpy as np
+
+start_token = 'B'
+end_token = 'E'
+
+
+def process_poem(file_name):
+    # 诗集
+    poem = []
+    with open(file_name, "r", encoding='utf-8', ) as f:
+        for line in f.readlines():
+            try:
+                title, content = line.strip().split(':')
+                content = content.replace(' ', '')
+                if '_' in content or '(' in content or '（' in content or '《' in content or '[' in content or \
+                        start_token in content or end_token in content:
+                    continue
+                if len(content) < 5 or len(content) > 79:
+                    continue
+                content = start_token + content + end_token
+                poem.append(content)
+            except ValueError as e:
+                pass
+    poem = sorted(poem, key=lambda l: len(line))
+
+    all_words = []
+    for poem in poem:
+        all_words += [word for word in poem]
+    counter = collections.Counter(all_words)
+    count_pairs = sorted(counter.items(), key=lambda x: -x[1])
+    words, _ = zip(*count_pairs)
+
+    words = words[:len(words)] + (' ',)
+    word_int_map = dict(zip(words, range(len(words))))
+    poem_vector = [list(map(lambda word: word_int_map.get(word, len(words)), poem)) for poem in poem]
+
+    return poem_vector, word_int_map, words
+
+
+def generate_batch(batch_size, poem_vec, word_to_int):
+    n_chunk = len(poem_vec) // batch_size
+    x_batches = []
+    y_batches = []
+    for i in range(n_chunk):
+        start_index = i * batch_size
+        end_index = start_index + batch_size
+
+        batches = poem_vec[start_index:end_index]
+        length = max(map(len, batches))
+        x_data = np.full((batch_size, length), word_to_int[' '], np.int32)
+        for row in range(batch_size):
+            x_data[row, :len(batches[row])] = batches[row]
+        y_data = np.copy(x_data)
+        y_data[:, :-1] = x_data[:, 1:]
+        """
+        x_data             y_data
+        [6,2,4,6,9]       [2,4,6,9,9]
+        [1,4,2,8,5]       [4,2,8,5,5]
+        """
+        x_batches.append(x_data)
+        y_batches.append(y_data)
+    return x_batches, y_batches

src/train.py

@@ -0,0 +1,68 @@
+import os
+import numpy as np
+import tensorflow as tf
+from poem.model import rnn_model
+from poem.poem import process_poem, generate_batch
+
+tf.app.flags.DEFINE_integer('batch_size', 64, 'batch size.')
+tf.app.flags.DEFINE_float('learning_rate', 0.01, 'learning rate.')
+tf.app.flags.DEFINE_string('model_dir', os.path.abspath('./model'), 'model save path.')
+tf.app.flags.DEFINE_string('file_path', os.path.abspath('./data/poem.txt'), 'file name of poem.')
+tf.app.flags.DEFINE_string('model_prefix', 'poem', 'model save prefix.')
+tf.app.flags.DEFINE_integer('epochs', 50, 'train how many epochs.')
+
+FLAGS = tf.app.flags.FLAGS
+
+
+def run_training():
+    if not os.path.exists(FLAGS.model_dir):
+        os.makedirs(FLAGS.model_dir)
+
+    poem_vector, word_to_int, vocabularies = process_poem(FLAGS.file_path)
+    batches_inputs, batches_outputs = generate_batch(FLAGS.batch_size, poem_vector, word_to_int)
+
+    input_data = tf.placeholder(tf.int32, [FLAGS.batch_size, None])
+    output_targets = tf.placeholder(tf.int32, [FLAGS.batch_size, None])
+
+    end_points = rnn_model(model='lstm', input_data=input_data, output_data=output_targets, vocab_size=len(
+        vocabularies), rnn_size=128, num_layers=2, batch_size=64, learning_rate=FLAGS.learning_rate)
+
+    saver = tf.train.Saver(tf.global_variables())
+    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
+    with tf.Session() as sess:
+        # sess = tf_debug.LocalCLIDebugWrapperSession(sess=sess)
+        # sess.add_tensor_filter("has_inf_or_nan", tf_debug.has_inf_or_nan)
+        sess.run(init_op)
+
+        start_epoch = 0
+        checkpoint = tf.train.latest_checkpoint(FLAGS.model_dir)
+        if checkpoint:
+            saver.restore(sess, checkpoint)
+            print("## restore from the checkpoint {0}".format(checkpoint))
+            start_epoch += int(checkpoint.split('-')[-1])
+        print('## start training...')
+        try:
+            for epoch in range(start_epoch, FLAGS.epochs):
+                n = 0
+                n_chunk = len(poem_vector) // FLAGS.batch_size
+                for batch in range(n_chunk):
+                    loss, _, _ = sess.run([
+                        end_points['total_loss'],
+                        end_points['last_state'],
+                        end_points['train_op']
+                    ], feed_dict={input_data: batches_inputs[n], output_targets: batches_outputs[n]})
+                    n += 1
+                    print('Epoch: %d, batch: %d, training loss: %.6f' % (epoch, batch, loss))
+                if epoch % 6 == 0:
+                    saver.save(sess, os.path.join(FLAGS.model_dir, FLAGS.model_prefix), global_step=epoch)
+        except KeyboardInterrupt:
+            print('## Interrupt manually, try saving checkpoint for now...')
+            saver.save(sess, os.path.join(FLAGS.model_dir, FLAGS.model_prefix), global_step=epoch)
+            print('## Last epoch were saved, next time will start from epoch {}.'.format(epoch))
+
+
+def main(_):
+    run_training()
+
+if __name__ == '__main__':
+    tf.app.run()

src/utils/clean_cn.py

@@ -0,0 +1,141 @@
+"""
+this script using for clean Chinese corpus.
+you can set level for clean, i.e.:
+level='all', will clean all character that not Chinese, include punctuations
+level='normal', this will generate corpus like normal use, reserve alphabets and numbers
+level='clean', this will remove all except Chinese and Chinese punctuations
+
+besides, if you want remove complex Chinese characters, just set this to be true:
+simple_only=True
+"""
+import numpy as np
+import os
+import string
+
+
+cn_punctuation_set = ['，', '。', '！', '？', '"', '"', '、']
+en_punctuation_set = [',', '.', '?', '!', '"', '"']
+
+
+def clean_cn_corpus(file_name, clean_level='all', simple_only=True, is_save=True):
+    """
+    clean Chinese corpus.
+    :param file_name:
+    :param clean_level:
+    :param simple_only:
+    :param is_save:
+    :return: clean corpus in list type.
+    """
+    if os.path.dirname(file_name):
+        base_dir = os.path.dirname(file_name)
+    else:
+        print('not set dir. please check')
+
+    save_file = os.path.join(base_dir, os.path.basename(file_name).split('.')[0] + '_cleaned.txt')
+    with open(file_name, 'r+') as f:
+        clean_content = []
+        for l in f.readlines():
+            l = l.strip()
+            if l == '':
+                pass
+            else:
+                l = list(l)
+                should_remove_words = []
+                for w in l:
+                    if not should_reserve(w, clean_level):
+                        should_remove_words.append(w)
+                clean_line = [c for c in l if c not in should_remove_words]
+                clean_line = ''.join(clean_line)
+                if clean_line != '':
+                    clean_content.append(clean_line)
+    if is_save:
+        with open(save_file, 'w+') as f:
+            for l in clean_content:
+                f.write(l + '\n')
+        print('[INFO] cleaned file have been saved to %s.' % save_file)
+    return clean_content
+
+
+def should_reserve(w, clean_level):
+    if w == ' ':
+        return True
+    else:
+        if clean_level == 'all':
+            # only reserve Chinese characters
+            if w in cn_punctuation_set or w in string.punctuation or is_alphabet(w):
+                return False
+            else:
+                return is_chinese(w)
+        elif clean_level == 'normal':
+            # reserve Chinese characters, English alphabet, number
+            if is_chinese(w) or is_alphabet(w) or is_number(w):
+                return True
+            elif w in cn_punctuation_set or w in en_punctuation_set:
+                return True
+            else:
+                return False
+        elif clean_level == 'clean':
+            if is_chinese(w):
+                return True
+            elif w in cn_punctuation_set:
+                return True
+            else:
+                return False
+        else:
+            raise "clean_level not support %s, please set for all, normal, clean" % clean_level
+
+
+def is_chinese(uchar):
+    """is chinese"""
+    if u'\u4e00' <= uchar <= u'\u9fa5':
+        return True
+    else:
+        return False
+
+
+def is_number(uchar):
+    """is number"""
+    if u'\u0030' <= uchar <= u'\u0039':
+        return True
+    else:
+        return False
+
+
+def is_alphabet(uchar):
+    """is alphabet"""
+    if (u'\u0041' <= uchar <= u'\u005a') or (u'\u0061' <= uchar <= u'\u007a'):
+        return True
+    else:
+        return False
+
+
+def semi_angle_to_sbc(uchar):
+    """半角转全角"""
+    inside_code = ord(uchar)
+    if inside_code < 0x0020 or inside_code > 0x7e:
+        return uchar
+    if inside_code == 0x0020:
+        inside_code = 0x3000
+    else:
+        inside_code += 0xfee0
+    return chr(inside_code)
+
+
+def sbc_to_semi_angle(uchar):
+    """全角转半角"""
+    inside_code = ord(uchar)
+    if inside_code == 0x3000:
+        inside_code = 0x0020
+    else:
+        inside_code -= 0xfee0
+    if inside_code < 0x0020 or inside_code > 0x7e:
+        return uchar
+    return chr(inside_code)
+
+
+
+
+
+
+
+

train.py

@@ -0,0 +1,69 @@
+import os
+import numpy as np
+import tensorflow as tf
+from poem.model import rnn_model
+from poem.poem import process_poem, generate_batch
+
+tf.app.flags.DEFINE_integer('batch_size', 64, 'batch size.')
+tf.app.flags.DEFINE_float('learning_rate', 0.01, 'learning rate.')
+tf.app.flags.DEFINE_string('model_dir', os.path.abspath('./model'), 'model save path.')
+tf.app.flags.DEFINE_string('file_path', os.path.abspath('./data/poem.txt'), 'file name of poem.')
+tf.app.flags.DEFINE_string('model_prefix', 'poem', 'model save prefix.')
+tf.app.flags.DEFINE_integer('epochs', 50, 'train how many epochs.')
+
+FLAGS = tf.app.flags.FLAGS
+
+
+def run_training():
+    if not os.path.exists(FLAGS.model_dir):
+        os.makedirs(FLAGS.model_dir)
+
+    poem_vector, word_to_int, vocabularies = process_poem(FLAGS.file_path)
+    batches_inputs, batches_outputs = generate_batch(FLAGS.batch_size, poem_vector, word_to_int)
+
+    input_data = tf.placeholder(tf.int32, [FLAGS.batch_size, None])
+    output_targets = tf.placeholder(tf.int32, [FLAGS.batch_size, None])
+
+    end_points = rnn_model(model='lstm', input_data=input_data, output_data=output_targets, vocab_size=len(
+        vocabularies), rnn_size=128, num_layers=2, batch_size=64, learning_rate=FLAGS.learning_rate)
+
+    saver = tf.train.Saver(tf.global_variables())
+    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
+    with tf.Session() as sess:
+        # sess = tf_debug.LocalCLIDebugWrapperSession(sess=sess)
+        # sess.add_tensor_filter("has_inf_or_nan", tf_debug.has_inf_or_nan)
+        sess.run(init_op)
+
+        start_epoch = 0
+        checkpoint = tf.train.latest_checkpoint(FLAGS.model_dir)
+        if checkpoint:
+            saver.restore(sess, checkpoint)
+            print("## restore from the checkpoint {0}".format(checkpoint))
+            start_epoch += int(checkpoint.split('-')[-1])
+        print('## start training...')
+        try:
+            for epoch in range(start_epoch, FLAGS.epochs):
+                n = 0
+                n_chunk = len(poem_vector) // FLAGS.batch_size
+                for batch in range(n_chunk):
+                    loss, _, _ = sess.run([
+                        end_points['total_loss'],
+                        end_points['last_state'],
+                        end_points['train_op']
+                    ], feed_dict={input_data: batches_inputs[n], output_targets: batches_outputs[n]})
+                    n += 1
+                    print('Epoch: %d, batch: %d, training loss: %.6f' % (epoch, batch, loss))
+                if epoch % 6 == 0:
+                    saver.save(sess, os.path.join(FLAGS.model_dir, FLAGS.model_prefix), global_step=epoch)
+        except KeyboardInterrupt:
+            print('## Interrupt manually, try saving checkpoint for now...')
+            saver.save(sess, os.path.join(FLAGS.model_dir, FLAGS.model_prefix), global_step=epoch)
+            print('## Last epoch were saved, next time will start from epoch {}.'.format(epoch))
+
+
+def main(_):
+    run_training()
+
+
+if __name__ == '__main__':
+    tf.app.run()