lstm2.py

"""
Usage:
    lstm_text_generation.py <DIR>

Options:
    <DIR>      Directory containing parsed subtitles (e.g., data/lacollectioneuse.txt)
"""
from __future__ import print_function
from keras.models import Sequential, model_from_json
from keras.optimizers import Adagrad
from keras.layers.core import Dense, Activation, Dropout
from keras.layers.recurrent import LSTM
# from keras.datasets.data_utils import get_file
from keras.utils.data_utils import get_file

from os import path, symlink
import numpy as np
import random, sys
import re
from glob import glob

from docopt import docopt
args = docopt(__doc__)

'''
    Example script to generate text from Nietzsche's writings.

    At least 20 epochs are required before the generated text
    starts sounding coherent.

    It is recommended to run this script on GPU, as recurrent
    networks are quite computationally intensive.

    If you try this script on new data, make sure your corpus
    has at least ~100k characters. ~1M is better.
'''

_text = ''
for fname in glob(path.join(args["<DIR>"], '*.txt')):
    print("Reading file: %s" % fname)
    fh = open(fname, 'ro')
    _text += fh.read().decode('utf-8').lower()
    fh.close()
print('corpus length:', len(_text))

text = []
for w in re.finditer(r"(\w+)", _text, re.MULTILINE | re.DOTALL | re.UNICODE):
    text.append(w.group(1))
chars = set(text)
print('total chars:', len(chars))
char_indices = dict((c, i) for i, c in enumerate(chars))
indices_char = dict((i, c) for i, c in enumerate(chars))

# cut the text in semi-redundant sequences of maxlen characters
maxlen = 20
step = 3
sentences = []
next_chars = []
for i in range(0, len(text) - maxlen, step):
    sentences.append(text[i : i + maxlen])
    next_chars.append(text[i + maxlen])
print('nb sequences:', len(sentences))

print('Vectorization...')
X = np.zeros((len(sentences), maxlen, len(chars)), dtype=np.bool)
y = np.zeros((len(sentences), len(chars)), dtype=np.bool)
for i, sentence in enumerate(sentences):
    for t, char in enumerate(sentence):
        X[i, t, char_indices[char]] = 1
    y[i, char_indices[next_chars[i]]] = 1


# build the model: 2 stacked LSTM
try:
    print('Loading model...')
    fh = open('keras-manos.model.json', 'rb')
    model = model_from_json(fh.read())
    fh.close()
except Exception as error:
    print(error)
    print('Build model...')
    model = Sequential()
    model.add(LSTM(512, input_shape=(maxlen, len(chars)), activation='sigmoid', return_sequences=True))
    model.add(Dropout(0.2))
    model.add(LSTM(512, activation='sigmoid', return_sequences=False))
    model.add(Dropout(0.2))
    model.add(Dense(len(chars)))
    model.add(Activation('softmax'))

#    adagrad = Adagrad(lr=0.01, epsilon=1e-6, clipnorm=1.)
#    model.compile(loss='binary_crossentropy', optimizer=adagrad)
    model.compile(loss='categorical_crossentropy', optimizer='rmsprop')

    fh = open('keras-manos.model.json', 'wb')
    fh.write(model.to_json())
    fh.close()


# helper function to sample an index from a probability array
def sample(a, temperature=1.0):
    a = np.log(a) / temperature
    dist = np.exp(a)/np.sum(np.exp(a))
    choices = range(len(a))
    return np.random.choice(choices, p=dist)
    # a = np.log(a)/temperature
    # a = np.exp(a)/np.sum(np.exp(a))
    # return np.argmax(np.random.multinomial(1,a,1))

# Load previous weights
if path.isfile('keras-manos.weights.latest.h5'):
    model.load_weights('keras-manos.weights.latest.h5')

# train the model, output generated text after each iteration
for iteration in range(1, 3):
    print()
    print('-' * 50)
    print('Iteration', iteration)
    model.fit(X, y, batch_size=128, epochs=1)


    print(text)
    start_index = random.randint(0, len(text) - maxlen - 1)

    for diversity in [0.2, 0.5, 1.0, 1.2]:
        print()
        print('----- diversity:', diversity)

        generated = ''
        sentence = text[start_index : start_index + maxlen]
        generated += " ".join(sentence).encode('utf-8')
        print('----- Generating with seed: "' + " ".join(sentence).encode('utf-8') + '"')
        sys.stdout.write(generated)

        for _ in range(400):
            x = np.zeros((1, maxlen, len(chars)))
            for t, char in enumerate(sentence):
                x[0, t, char_indices[char]] = 1.

            preds = model.predict(x, verbose=0)[0]
            next_index = sample(preds, diversity)
            next_char = indices_char[next_index]

            # generated += next_char
            # sentence = sentence[1:] + next_char

            sys.stdout.write(' ' + next_char.encode('utf-8'))
            sys.stdout.flush()
        print()

    model.save_weights('keras-manos.weights.iter-%d.h5' % iteration)
    symlink('keras-manos.weights.iter-%d.h5' % iteration, 'keras-manos.weights.latest.h5')