cadl.charrnn — CADL 1.0.0 documentation

# Source code for cadl.charrnn

    """Character-level Recurrent Neural Network.
    """
    """
    Copyright 2017 Parag K. Mital.  See also NOTICE.md.

    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

        http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
    """

    import tensorflow as tf
    import numpy as np
    import os
    import sys
    import collections
    import gzip
    from cadl import utils

    def build_model(txt,
                    batch_size=1,
                    sequence_length=1,
                    n_layers=2,
                    n_cells=100,
                    gradient_clip=10.0,
                    learning_rate=0.001):
        """Summary

        Parameters
        ----------
        txt : TYPE
            Description
        batch_size : int, optional
            Description
        sequence_length : int, optional
            Description
        n_layers : int, optional
            Description
        n_cells : int, optional
            Description
        gradient_clip : float, optional
            Description
        learning_rate : float, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        vocab = list(set(txt))
        vocab.sort()
        n_chars = len(vocab)
        encoder = collections.OrderedDict(zip(vocab, range(n_chars)))
        decoder = collections.OrderedDict(zip(range(n_chars), vocab))

        X = tf.placeholder(tf.int32, [None, sequence_length], name='X')
        Y = tf.placeholder(tf.int32, [None, sequence_length], name='Y')
        keep_prob = tf.placeholder(tf.float32, name='keep_prob')

        with tf.variable_scope('embedding'):
            embedding = tf.get_variable("embedding", [n_chars, n_cells])
            # Each sequence element will be connected to n_cells
            Xs = tf.nn.embedding_lookup(embedding, X)
            # Then slice each sequence element, giving us sequence number of
            # batch x 1 x n_chars Tensors
            Xs = tf.split(axis=1, num_or_size_splits=sequence_length, value=Xs)
            # Get rid of singleton sequence element dimension
            Xs = [tf.squeeze(X_i, [1]) for X_i in Xs]

        with tf.variable_scope('rnn'):
            cells = tf.contrib.rnn.MultiRNNCell([
                tf.contrib.rnn.DropoutWrapper(
                    tf.contrib.rnn.BasicLSTMCell(
                        num_units=n_cells, forget_bias=0.0, state_is_tuple=True),
                    output_keep_prob=keep_prob) for _ in range(n_layers)
            ])
            initial_state = cells.zero_state(tf.shape(X)[0], tf.float32)
            # returns a length sequence length list of outputs, one for each input
            outputs, final_state = tf.contrib.rnn.static_rnn(
                cells, Xs, initial_state=initial_state)
            # now concat the sequence length number of batch x n_cells Tensors to
            # give [sequence_length x batch, n_cells]
            outputs_flat = tf.reshape(
                tf.concat(axis=1, values=outputs), [-1, n_cells])

        with tf.variable_scope('prediction'):
            W = tf.get_variable(
                "W",
                shape=[n_cells, n_chars],
                initializer=tf.contrib.layers.xavier_initializer())
            b = tf.get_variable(
                "b", shape=[n_chars], initializer=tf.constant_initializer())
            logits = tf.matmul(outputs_flat, W) + b
            probs = tf.nn.softmax(logits)
            Y_pred = tf.argmax(probs, 1)

        with tf.variable_scope('loss'):
            loss = tf.contrib.legacy_seq2seq.sequence_loss_by_example([logits], [
                tf.reshape(tf.concat(axis=1, values=Y), [-1])
            ], [tf.ones([batch_size * sequence_length])])
            cost = tf.reduce_sum(loss) / batch_size

        with tf.name_scope('optimizer'):
            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
            gradients = []
            clip = tf.constant(gradient_clip, name="clip")
            for grad, var in optimizer.compute_gradients(cost):
                gradients.append((tf.clip_by_value(grad, -clip, clip), var))
            updates = optimizer.apply_gradients(gradients)

        model = {
            'X': X,
            'Y': Y,
            'logits': logits,
            'probs': probs,
            'Y_pred': Y_pred,
            'keep_prob': keep_prob,
            'cost': cost,
            'updates': updates,
            'initial_state': initial_state,
            'final_state': final_state,
            'decoder': decoder,
            'encoder': encoder,
            'vocab_size': n_chars
        }
        return model

    def train(txt,
              batch_size=100,
              sequence_length=150,
              n_cells=200,
              n_layers=3,
              learning_rate=0.00001,
              max_iter=50000,
              gradient_clip=5.0,
              ckpt_name="model.ckpt",
              keep_prob=1.0):
        """train

        Parameters
        ----------
        txt : TYPE
            Description
        batch_size : int, optional
            Description
        sequence_length : int, optional
            Description
        n_cells : int, optional
            Description
        n_layers : int, optional
            Description
        learning_rate : float, optional
            Description
        max_iter : int, optional
            Description
        gradient_clip : float, optional
            Description
        ckpt_name : str, optional
            Description
        keep_prob : float, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        g = tf.Graph()
        with tf.Session(graph=g) as sess:
            model = build_model(
                txt=txt,
                batch_size=batch_size,
                sequence_length=sequence_length,
                n_layers=n_layers,
                n_cells=n_cells,
                gradient_clip=gradient_clip,
                learning_rate=learning_rate)

            init_op = tf.group(tf.global_variables_initializer(),
                               tf.local_variables_initializer())
            saver = tf.train.Saver()
            sess.run(init_op)
            if os.path.exists(ckpt_name + '.index') or os.path.exists(ckpt_name):
                saver.restore(sess, ckpt_name)
                print("Model restored.")

            cursor = 0
            it_i = 0
            print_step = 1000
            avg_cost = 0
            while it_i < max_iter:
                Xs, Ys = [], []
                for batch_i in range(batch_size):
                    Xs.append([
                        model['encoder'][ch]
                        for ch in txt[cursor:cursor + sequence_length]
                    ])
                    Ys.append([
                        model['encoder'][ch]
                        for ch in txt[cursor + 1:cursor + sequence_length + 1]
                    ])
                    cursor += sequence_length
                    if (cursor + 1) >= len(txt) - sequence_length - 1:
                        cursor = np.random.randint(0, high=sequence_length)

                feed_dict = {
                    model['X']: Xs,
                    model['Y']: Ys,
                    model['keep_prob']: keep_prob
                }
                out = sess.run(
                    [model['cost'], model['updates']], feed_dict=feed_dict)
                avg_cost += out[0]

                if (it_i + 1) % print_step == 0:
                    p = sess.run(
                        model['probs'],
                        feed_dict={
                            model['X']: np.array(Xs[-1])[np.newaxis],
                            model['keep_prob']: 1.0
                        })
                    print(p.shape, 'min:',
                          np.min(p), 'max:',
                          np.max(p), 'mean:', np.mean(p), 'std:', np.std(p))
                    if isinstance(txt[0], str):
                        # Print original string
                        print('original:',
                              "".join([model['decoder'][ch] for ch in Xs[-1]]))

                        # Print max guess
                        amax = []
                        for p_i in p:
                            amax.append(model['decoder'][np.argmax(p_i)])
                        print('synth(amax):', "".join(amax))

                        # Print w/ sampling
                        samp = []
                        for p_i in p:
                            p_i = p_i.astype(np.float64)
                            p_i = p_i / p_i.sum()
                            idx = np.argmax(np.random.multinomial(1, p_i.ravel()))
                            samp.append(model['decoder'][idx])
                        print('synth(samp):', "".join(samp))

                    print(it_i, avg_cost / print_step)
                    avg_cost = 0

                    save_path = saver.save(sess, ckpt_name, global_step=it_i)
                    print("Model saved in file: %s" % save_path)

                print(it_i, out[0], end='\r')
                it_i += 1

            return model

    def infer(txt,
              ckpt_name,
              n_iterations,
              n_cells=200,
              n_layers=3,
              learning_rate=0.001,
              max_iter=5000,
              gradient_clip=10.0,
              init_value=[0],
              keep_prob=1.0,
              sampling='prob',
              temperature=1.0):
        """infer

        Parameters
        ----------
        txt : TYPE
            Description
        ckpt_name : TYPE
            Description
        n_iterations : TYPE
            Description
        n_cells : int, optional
            Description
        n_layers : int, optional
            Description
        learning_rate : float, optional
            Description
        max_iter : int, optional
            Description
        gradient_clip : float, optional
            Description
        init_value : list, optional
            Description
        keep_prob : float, optional
            Description
        sampling : str, optional
            Description
        temperature : float, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        g = tf.Graph()
        with tf.Session(graph=g) as sess:
            sequence_length = len(init_value)
            model = build_model(
                txt=txt,
                batch_size=1,
                sequence_length=sequence_length,
                n_layers=n_layers,
                n_cells=n_cells,
                gradient_clip=gradient_clip,
                learning_rate=learning_rate)

            init_op = tf.group(tf.global_variables_initializer(),
                               tf.local_variables_initializer())
            saver = tf.train.Saver()
            sess.run(init_op)
            if os.path.exists(ckpt_name):
                saver.restore(sess, ckpt_name)
                print("Model restored.")

            state = []
            synth = [init_value]
            for s_i in model['final_state']:
                state += sess.run(
                    [s_i.c, s_i.h],
                    feed_dict={
                        model['X']: [synth[-1]],
                        model['keep_prob']: keep_prob
                    })

            for i in range(n_iterations):
                # print('iteration: {}/{}'.format(i, n_iterations), end='\r')
                feed_dict = {model['X']: [synth[-1]], model['keep_prob']: keep_prob}
                state_updates = []
                for state_i in range(n_layers):
                    feed_dict[model['initial_state'][state_i].c] = \
                        state[state_i * 2]
                    feed_dict[model['initial_state'][state_i].h] = state[state_i * 2
                                                                         + 1]
                    state_updates.append(model['final_state'][state_i].c)
                    state_updates.append(model['final_state'][state_i].h)
                p = sess.run(model['probs'], feed_dict=feed_dict)[0]
                if sampling == 'max':
                    p = np.argmax(p)
                else:
                    p = p.astype(np.float64)
                    p = np.log(p) / temperature
                    p = np.exp(p) / np.sum(np.exp(p))
                    p = np.random.multinomial(1, p.ravel())
                    p = np.argmax(p)
                # Get the current state
                state = [
                    sess.run(s_i, feed_dict=feed_dict) for s_i in state_updates
                ]
                synth.append([p])
                print(model['decoder'][p], end='')
                sys.stdout.flush()
                if model['decoder'][p] in ['.', '?', '!']:
                    print('\n')
            print(np.concatenate(synth).shape)
        print("".join([model['decoder'][ch] for ch in np.concatenate(synth)]))
        return [model['decoder'][ch] for ch in np.concatenate(synth)]

    def test_alice(max_iter=5):
        """Summary

        Parameters
        ----------
        max_iter : int, optional
            Description

        Returns
        -------
        TYPE
            Description
        """
        utils.download('https://s3.amazonaws.com/cadl/models/alice.txt.gz')
        with gzip.open('alice.txt.gz', 'rb') as fp:
            txt = fp.read().decode('utf-8')
        return train(txt, n_layers=2, n_cells=20, max_iter=max_iter)

    def test_trump(max_iter=100):
        """Summary

        Parameters
        ----------
        max_iter : int, optional
            Description
        """
        utils.download(
            'https://s3.amazonaws.com/cadl/models/trump.data-00000-of-00001')
        utils.download('https://s3.amazonaws.com/cadl/models/trump.meta')
        utils.download('https://s3.amazonaws.com/cadl/models/trump.index')
        utils.download('https://s3.amazonaws.com/cadl/models/trump.txt')
        with open('trump.txt', 'r') as fp:
            txt = fp.read()
        #train(txt, ckpt_name='trump', max_iter=max_iter)
        print(infer(txt, ckpt_name='./trump', n_iterations=max_iter))

    def test_wtc():
        """Summary
        """
        from scipy.io.wavfile import write, read
        rate, aud = read('wtc.wav')
        txt = np.int8(np.round(aud / 16384.0 * 128.0))
        txt = np.squeeze(txt).tolist()
        # try with more than 100 iterations, e.g. 50k - 200k
        train(txt, sequence_length=250, n_layers=3, n_cells=512, max_iter=100)
        synthesis = infer(
            txt,
            './model.ckpt',
            8000 * 30,
            n_layers=3,
            n_cells=150,
            keep_prob=1.0,
            sampling='prob')
        snd = np.int16(np.array(synthesis) / 128.0 * 16384.0)
        write('wtc-synth.wav', 8000, snd)

    if __name__ == '__main__':
        test_alice()