# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

# %matplotlib inline
import utils; reload(utils)
from utils import *

# ### Setup

text = open('../data/nietzsche.txt').read();

chars = sorted(list(set(text)))
vocab_size = len(chars) + 1

chars.insert(0, '\0')

''.join(chars)[1:-6]

char_indices = dict((c, i) for i, c in enumerate(chars))
indices_char = dict((i, c) for i, c in enumerate(chars))

idx = [char_indices[c] for c in text]

idx[:10]

''.join([indices_char[i] for i in idx[:50]])

# ### 3 char model

cs=3
c1_dat = [idx[i] for i in xrange(0, len(idx)-1-cs, cs)]
c2_dat = [idx[i+1] for i in xrange(0, len(idx)-1-cs, cs)]
c3_dat = [idx[i+2] for i in xrange(0, len(idx)-1-cs, cs)]
c4_dat = [idx[i+3] for i in xrange(0, len(idx)-1-cs, cs)]

x1 = np.stack(c1_dat[:-2])
x2 = np.stack(c2_dat[:-2])
x3 = np.stack(c3_dat[:-2])

y = np.stack(c4_dat[:-2])

x1[:4], x2[:4], x3[:4]

y[:4]

x1.shape, y.shape

n_fac = 42


def embedding_input(name, n_in, n_out):
    inp = Input(shape=(1,), dtype='int64', name=name)
    emb = Embedding(n_in, n_out, input_length=1)(inp)
    return inp, Flatten()(emb)


c1_in, c1 = embedding_input('c1', vocab_size, n_fac)
c2_in, c2 = embedding_input('c2', vocab_size, n_fac)
c3_in, c3 = embedding_input('c3', vocab_size, n_fac)

# ### Create and train model

n_hidden = 256

dense_in = Dense(n_hidden, activation='relu')

c1_hidden = dense_in(c1)

dense_hidden = Dense(n_hidden, activation='relu')

c2_dense = dense_in(c2)
hidden_2 = dense_hidden(c1_hidden)
c2_hidden = merge([c2_dense, hidden_2])

c3_dense = dense_in(c3)
hidden_3 = dense_hidden(c2_hidden)
c3_hidden = merge([c3_dense, hidden_3])

dense_out = Dense(vocab_size, activation='softmax')

c4_out = dense_out(c3_hidden)

model = Model([c1_in, c2_in, c3_in], c4_out)

model.compile(loss='sparse_categorical_crossentropy', optimizer=Adam())

model.optimizer.lr = 0.00001

model.fit([x1, x2, x3], y, batch_size=64, nb_epoch=4)


# ### Test model

def get_next(inp):
    idxs = [char_indices[c] for c in inp]
    arrs = [np.array(i)[np.newaxis] for i in idxs]
    p = model.predict(arrs)
    i = np.argmax(p)
    print(p)
    print(i)
    return chars[i]


get_next(' an')

# **Holy shit, this does not work!**

# ### Our first RNN!

# #### Create inputs

cs=8

c_in_dat = [ [idx[i+n] for i in xrange(0, len(idx)-1-cs, cs)] for n in range(cs) ]

print(idx[:40])

for i in range(len(c_in_dat)):
    print(c_in_dat[i][:20])

c_out_dat = [idx[i+cs] for i in xrange(0, len(idx)-1-cs, cs)]

xs = [np.stack(c[:-2]) for c in c_in_dat]

len(xs), xs[0].shape

y = np.stack(c_out_dat[:-2])

[xs[n][:cs] for n in range(cs)]

y[:cs]

n_fac = 42


# ### Create and train model

def embedding_input(name, n_in, n_out):
    inp = Input(shape=(1,), dtype='int64', name=name+'_in')
    emb = Embedding(n_in, n_out, input_length=1, name=name+'_emb')(inp)
    return inp, Flatten()(emb)


c_ins = [embedding_input('c'+str(n), vocab_size, n_fac) for n in range(cs)]

n_hidden = 256

dense_in = Dense(n_hidden, activation='relu', name="dense_in")
dense_hidden = Dense(n_hidden, activation='relu', init='identity', name="dense_hidden")
dense_out = Dense(vocab_size, activation='softmax', name="dense_out")

hidden = dense_in(c_ins[0][1])

for i in range(1,cs):
    c_dense = dense_in(c_ins[i][1])
    hidden = dense_hidden(hidden)
    hidden = merge([c_dense, hidden])

c_out = dense_out(hidden)

model = Model([c[0] for c in c_ins], c_out)
model.compile(loss='sparse_categorical_crossentropy', optimizer=Adam())

from keras.utils import visualize_util
visualize_util.plot(model, to_file='model.png')
# plot_model(model, to_file='model.png')

model.fit(xs, y, batch_size=64, nb_epoch=5)

model.fit(xs, y, batch_size=64, nb_epoch=5)


# ### Test model

def get_next(inp):
    idxs = [np.array(char_indices[c])[np.newaxis] for c in inp]
    p = model.predict(idxs)
    return chars[np.argmax(p)]


get_next('for thos')

get_next('part of ')

get_next('queens a')

get_next('there is')

# ### Our first RNN with Keras

n_hidden, n_fac, cs, vocab_size = (256, 42, 8, 86)

model = Sequential([
    Embedding(input_dim=vocab_size, output_dim=n_fac, input_length=cs),
    SimpleRNN(output_dim=n_hidden, activation='relu', inner_init='identity'), # what does identity mean?
    Dense(output_dim=vocab_size, activation='softmax')
])

model.summary()

model.compile(optimizer=Adam(), loss='sparse_categorical_crossentropy')

model.fit(np.concatenate(xs, axis=1), y, batch_size=64, nb_epoch=4)

model.fit(np.concatenate(xs, axis=1), y, batch_size=64, nb_epoch=4)


def get_next(inp):
    idxs = [np.array(char_indices[c])[np.newaxis] for c in inp]
    p = model.predict(idxs)
    return chars[np.argmax(p)]

StarcoderdataJupyterScriptsDedupFiltered