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

# + [markdown] colab_type="text" id="ovpZyIhNIgoq"
# ___
#
# <a href='http://www.pieriandata.com'> <img src='../Pierian_Data_Logo.png' /></a>
# ___
# # Text Generation with Neural Networks
#
# In this notebook we will create a network that can generate text, here we show it being done character by character. Very awesome write up on this here: http://karpathy.github.io/2015/05/21/rnn-effectiveness/
#
# We organized the process into "steps" so you can easily follow along with your own data sets.
# -

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tensorflow as tf

# +
# IGNORE THE CONTENT OF THIS CELL
# tf.compat.v1.disable_eager_execution()
# -

# ## Step 1: The Data
#
# You can grab any free text you want from here: https://www.gutenberg.org/
#
# We'll choose all of shakespeare's works (which we have already downloaded for you), mainly for two reasons:
#
# 1. Its a large corpus of text, its usually recommended you have at least a source of 1 million characters total to get realistic text generation.
#
# 2. It has a very distinctive style. Since the text data uses old style english and is formatted in the style of a stage play, it will be very obvious to us if the model is able to reproduce similar results.

# + colab={} colab_type="code" id="pD_55cOxLkAb"
path_to_file = 'shakespeare.txt'

# + colab={} colab_type="code" id="aavnuByVymwK"
text = open(path_to_file, 'r').read()

# + colab={} colab_type="code" id="Duhg9NrUymwO"
print(text[:500])
# -

# ### Understanding unique characters

# + colab={} colab_type="code" id="IlCgQBRVymwR"
# The unique characters in the file
vocab = sorted(set(text))
print(vocab)
len(vocab)

# + [markdown] colab_type="text" id="rNnrKn_lL-IJ"
# ## Step 2: Text Processing

# + [markdown] colab_type="text" id="LFjSVAlWzf-N"
# ### Text Vectorization
#
# We know a neural network can't take in the raw string data, we need to assign numbers to each character. Let's create two dictionaries that can go from numeric index to character and character to numeric index.

# + colab={} colab_type="code" id="IalZLbvOzf-F"
char_to_ind = {u:i for i, u in enumerate(vocab)}
# -

char_to_ind

# + colab={} colab_type="code" id="IalZLbvOzf-F"
ind_to_char = np.array(vocab)
# -

ind_to_char

# + colab={} colab_type="code" id="IalZLbvOzf-F"
encoded_text = np.array([char_to_ind[c] for c in text])
# -

encoded_text

# + [markdown] colab_type="text" id="tZfqhkYCymwX"
# We now have a mapping we can use to go back and forth from characters to numerics.
# -

sample = text[:20]
sample

encoded_text[:20]

# + [markdown] colab_type="text" id="bbmsf23Bymwe"
# ## Step 3: Creating Batches
#
# Overall what we are trying to achieve is to have the model predict the next highest probability character given a historical sequence of characters. Its up to us (the user) to choose how long that historic sequence. Too short a sequence and we don't have enough information (e.g. given the letter "a" , what is the next character) , too long a sequence and training will take too long and most likely overfit to sequence characters that are irrelevant to characters farther out. While there is no correct sequence length choice, you should consider the text itself, how long normal phrases are in it, and a reasonable idea of what characters/words are relevant to each other.
# -

print(text[:500])

line = "From fairest creatures we desire increase"

len(line)

part_stanza = """From fairest creatures we desire increase,
  That thereby beauty's rose might never die,
  But as the riper should by time decease,"""

len(part_stanza)

# + [markdown] colab_type="text" id="hgsVvVxnymwf"
# ### Training Sequences
#
# The actual text data will be the text sequence shifted one character forward. For example:
#
# Sequence In: "<NAME>"
# Sequence Out: "<NAME>"
#
#
# We can use the `tf.data.Dataset.from_tensor_slices` function to convert a text vector into a stream of character indices.

# + colab={} colab_type="code" id="0UHJDA39zf-O"
seq_len = 120
# -

total_num_seq = len(text)//(seq_len+1)

total_num_seq

# + colab={} colab_type="code" id="0UHJDA39zf-O"
# Create Training Sequences
char_dataset = tf.data.Dataset.from_tensor_slices(encoded_text)

for i in char_dataset.take(500):
     print(ind_to_char[i.numpy()])

# + [markdown] colab_type="text" id="-ZSYAcQV8OGP"
# The **batch** method converts these individual character calls into sequences we can feed in as a batch. We use seq_len+1 because of zero indexing. Here is what drop_remainder means:
#
# drop_remainder: (Optional.) A `tf.bool` scalar `tf.Tensor`, representing
#     whether the last batch should be dropped in the case it has fewer than
#     `batch_size` elements; the default behavior is not to drop the smaller
#     batch.
#

# + colab={} colab_type="code" id="l4hkDU3i7ozi"
sequences = char_dataset.batch(seq_len+1, drop_remainder=True)


# + [markdown] colab_type="text" id="UbLcIPBj_mWZ"
# Now that we have our sequences, we will perform the following steps for each one to create our target text sequences:
#
# 1. Grab the input text sequence
# 2. Assign the target text sequence as the input text sequence shifted by one step forward
# 3. Group them together as a tuple

# + colab={} colab_type="code" id="9NGu-FkO_kYU"
def create_seq_targets(seq):
    input_txt = seq[:-1]
    target_txt = seq[1:]
    return input_txt, target_txt


# + colab={} colab_type="code" id="9NGu-FkO_kYU"
dataset = sequences.map(create_seq_targets)
# -

for input_txt, target_txt in  dataset.take(1):
    print(input_txt.numpy())
    print(''.join(ind_to_char[input_txt.numpy()]))
    print('\n')
    print(target_txt.numpy())
    # There is an extra whitespace!
    print(''.join(ind_to_char[target_txt.numpy()]))

# + [markdown] colab_type="text" id="MJdfPmdqzf-R"
# ### Generating training batches
#
# Now that we have the actual sequences, we will create the batches, we want to shuffle these sequences into a random order, so the model doesn't overfit to any section of the text, but can instead generate characters given any seed text.

# + colab={} colab_type="code" id="p2pGotuNzf-S"
# Batch size
batch_size = 128

# Buffer size to shuffle the dataset so it doesn't attempt to shuffle
# the entire sequence in memory. Instead, it maintains a buffer in which it shuffles elements
buffer_size = 10000

dataset = dataset.shuffle(buffer_size).batch(batch_size, drop_remainder=True)

# + colab={} colab_type="code" id="p2pGotuNzf-S"
dataset

# + [markdown] colab_type="text" id="r6oUuElIMgVx"
# ## Step 4: Creating the Model

# + [markdown] colab_type="text" id="m8gPwEjRzf-Z"
# We will use an LSTM based model with a few extra features, including an embedding layer to start off with and **two** LSTM layers. We based this model architecture off the [DeepMoji](https://deepmoji.mit.edu/) and the original source code can be found [here](https://github.com/bfelbo/DeepMoji).
#
# The embedding layer will serve as the input layer, which essentially creates a lookup table that maps the numbers indices of each character to a vector with "embedding dim" number of dimensions. As you can imagine, the larger this embedding size, the more complex the training. This is similar to the idea behind word2vec, where words are mapped to some n-dimensional space. Embedding before feeding straight into the LSTM usually leads to more realisitic results.

# + colab={} colab_type="code" id="zHT8cLh7EAsg"
# Length of the vocabulary in chars
vocab_size = len(vocab)

# The embedding dimension
embed_dim = 256

# Number of RNN units
rnn_neurons = 256
# -

# Now let's create a function that easily adapts to different variables as shown above.

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM,Dense,Embedding,Dropout


# + colab={} colab_type="code" id="MtCrdfzEI2N0"
def create_model(vocab_size, embed_dim, rnn_neurons, batch_size):
    model = Sequential()
    model.add(Embedding(vocab_size,embed_dim,batch_input_shape=[batch_size,None]))
    model.add(LSTM(rnn_neurons,return_sequences=True,stateful=True))
    model.add(Dropout(0.2))
    model.add(LSTM(rnn_neurons,return_sequences=True,stateful=True))
    model.add(Dropout(0.2))
    model.compile(optimizer='adam', loss="sparse_categorical_crossentropy")
    return model


# + colab={} colab_type="code" id="wwsrpOik5zhv"
model = create_model(
  vocab_size = vocab_size,
  embed_dim=embed_dim,
  rnn_neurons=rnn_neurons,
  batch_size=batch_size)
# -

model.summary()

# + [markdown] colab_type="text" id="LJL0Q0YPY6Ee"
# ## Step 5: Training the model
#
# This will take way too long on just CPU, go use Google Collab for GPU.
# -

epochs = 10


# +
# model.fit(dataset,epochs=epochs)

# + [markdown] colab_type="text" id="kKkD5M6eoSiN"
# ## Step 6: Generating text
#
# We need to create a loop very similar to how we generated forecasts with RNN with Time Series data. Let's take a look:

# + colab={} colab_type="code" id="WvuwZBX5Ogfd"
def generate_text(model, start_string):
  # Evaluation step (generating text using the learned model)

  # Number of characters to generate
  num_generate = 1000

  # Converting our start string to numbers (vectorizing)
  input_eval = [char2idx[s] for s in start_string]
  input_eval = tf.expand_dims(input_eval, 0)

  # Empty string to store our results
  text_generated = []

  # Low temperatures results in more predictable text.
  # Higher temperatures results in more surprising text.
  # Experiment to find the best setting.
  temperature = 1.0

  # Here batch size == 1
  model.reset_states()
  for i in range(num_generate):
      predictions = model(input_eval)
      # remove the batch dimension
      predictions = tf.squeeze(predictions, 0)

      # using a categorical distribution to predict the word returned by the model
      predictions = predictions / temperature
      predicted_id = tf.random.categorical(predictions, num_samples=1)[-1,0].numpy()

      # We pass the predicted word as the next input to the model
      # along with the previous hidden state
      input_eval = tf.expand_dims([predicted_id], 0)

      text_generated.append(idx2char[predicted_id])

  return (start_string + ''.join(text_generated))


# + colab={} colab_type="code" id="ktovv0RFhrkn"
print(generate_text(model, start_string=u"ROMEO: "))

StarcoderdataJupyterScriptsDedupFiltered