TV Script Generation

load_dotenv()

FOLDER_PATH = Path("../../../files/posts/nano/tv-script-generation/"
                   "tv-script-generation/")
if not FOLDER_PATH.is_dir():
    FOLDER_PATH.mkdir(parents=True)

Embed = partial(EmbedBokeh, folder_path=FOLDER_PATH)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
assert device.type == "cuda", 'No GPU found. Please use a GPU to train your neural network.'
print("Using {}".format(device))

WordIndices = Collection[int]

class Scripts:
    """Seinfeld Scripts

    Args:
     environment_key: environment variable with the source location
     dialog_only: remove descriptive columns
    """
    def __init__(self, environment_key: str="SCRIPTS", dialog_only: bool=True) -> None:
        self.environment_key = environment_key
        self.dialog_only = dialog_only
        self._script_blob = None
        self._path = None
        self._lines = None
        self._tokens = None
        self._line_tokens = None
        return

    @property
    def path(self) -> Path:
        """The path to the file"""
        if self._path is None:
            load_dotenv(".env")
            self._path = Path(os.environ.get("SCRIPTS")).expanduser()
            assert self._path.is_file()
        return self._path

    @property
    def script_blob(self) -> str:
        """The input file as a string"""
        if self._script_blob is None:
            with open(self.path) as reader:
                self._script_blob = reader.read()
        return self._script_blob

    @property
    def line_tokens(self) -> list:
        """list of tokens for each line"""
        if self._line_tokens is None:
            self._line_tokens = [line.split(" ") for line in self.lines]
        return self._line_tokens

    @property
    def lines(self) -> list:
        """The lines of the script"""
        if self._lines is None:
            lines = self.script_blob.split("\n")
            if self.dialog_only:
                lines = lines[1:]
                lines = [(",").join(line.split(",")[2:-3]) for line in lines]
            self._lines = lines
        return self._lines

    @property
    def tokens(self) -> Counter:
        """The tokens and their counts"""
        if self._tokens is None:
            self._tokens = Counter()
            for token in self.script_blob.split():
                self._tokens[token] += 1
        return self._tokens

class ScriptInspector:
    """gets some basic counts

    Args:
     scripts: object with the scripts
    """
    def __init__(self, scripts: Scripts=None) -> None:
        self._scripts = scripts
        self._line_count = None
        self._count_per_line = None
        self._mean_words_per_line = None
        self._median_words_per_line = None
        self._max_words_per_line = None
        self._min_words_per_line = None
        self._token_count = None
        return

    @property
    def scripts(self) -> Scripts:
        """The scripts object"""
        if self._scripts is None:
            self._scripts = Scripts()
        return self._scripts

    @property
    def line_count(self) -> int:
        """Number of lines in the source"""
        if self._line_count is None:
            self._line_count = len(self.scripts.lines)
        return self._line_count

    @property
    def count_per_line(self) -> list:
        """tokens per line"""
        if self._count_per_line is None:
            self._count_per_line = [len(tokens)
                                    for tokens in self.scripts.line_tokens]
        return self._count_per_line

    @property
    def mean_words_per_line(self) -> float:
        """Average number of words per line"""
        if self._mean_words_per_line is None:
            self._mean_words_per_line = (sum(self.count_per_line)
                                         /self.line_count)
        return self._mean_words_per_line

    @property
    def median_words_per_line(self) -> float:
        """Median words per line in the scripts"""
        if self._median_words_per_line is None:
            self._median_words_per_line = numpy.median(self.count_per_line)
        return self._median_words_per_line

    @property
    def max_words_per_line(self) -> int:
        """Count of words in longest line"""
        if self._max_words_per_line is None:
            self._max_words_per_line = max(self.count_per_line)
        return self._max_words_per_line

    @property
    def min_words_per_line(self) -> int:
        """Count of words in shortest line"""
        if self._min_words_per_line is None:
            self._min_words_per_line = min(self.count_per_line)
        return self._min_words_per_line

    @property
    def token_count(self) -> int:
        """Number of tokens in the text"""
        if self._token_count is None:
            self._token_count = sum(self.scripts.tokens.values())
        return self._token_count

    def most_common_tokens(self, count: int=10) -> list:
        """token, count tuples in descending rank

        Args:
         count: number of tuples to return in the list
        """
        if count > 0:
            return self.scripts.tokens.most_common(count)
        return self.scripts.tokens.most_common()[count:]

    def line_range(self, start: int=0, stop: int=10) -> list:
        """lines within range

        Args:
         start: index of first line
         stop: upper bound for last line
        """
        return self.scripts.lines[start:stop]

inspector = ScriptInspector()

view_line_range = (0, 10)

words_per_line = pandas.DataFrame(inspector.count_per_line,
                                  columns=["line_counts"])
print(words_per_line.shape)

lines = (("Number of unique tokens", "{:,}".format(inspector.token_count)),
         ("Number of lines", "{:,}".format(inspector.line_count)),
         ("Words in longest line", "{:,}".format(inspector.max_words_per_line)),
         ("Average number of words in each line", "{:.2f}".format(
             inspector.mean_words_per_line)),
         ("Median Words Per Line", "{:.2f}".format(
             inspector.median_words_per_line)),
         ("Words in shortest line", "{}".format(inspector.min_words_per_line))
)
print(tabulate(lines, headers="Statistic Value".split(), tablefmt="orgtbl"))

index = words_per_line.line_counts.idxmax()
print(inspector.count_per_line[index])
print(inspector.scripts.lines[index])

print(inspector.scripts.lines[words_per_line.line_counts.idxmin()])

plot = words_per_line.line_counts.hvplot.kde(title="Word Counts Per Line Distribution")
plotter = plot.opts(width=600, height=600, tools=["hover"])
Embed(plotter, "line_counts.js")()

plot = words_per_line.line_counts.hvplot.box(title="Words Per Line")
plot = plot.opts(tools=["hover"])
Embed(plot, "line_counts_boxplot.js")()

lines = ((token, "{:,}".format(count))
         for token, count in inspector.most_common_tokens())
print(tabulate(lines,
               tablefmt="orgtbl", headers=["Token", "Count"]))

words, counts = zip(*inspector.most_common_tokens(20))
top_twenty = pandas.DataFrame([counts], columns=words).T.reset_index()
top_twenty.columns = ["Word", "Count"]
layout = top_twenty.hvplot.bar(x="Word", y="Count",
                               title="Twenty Most Used Words",
                               colormap="Category20")
layout.opts(height=500, width=600)
Embed(layout, "top_twenty.js")()

for line in inspector.line_range(stop=5):
    print(line)

def create_lookup_tables(text: list) -> tuple:
    """
    Create lookup tables for vocabulary

    Args:
     text The text of tv scripts split into words

    Returns: 
     A tuple of dicts (vocab_to_int, int_to_vocab)
    """
    text = set(text)
    vocabulary_to_index = {token: index for index, token in enumerate(text)}
    index_to_vocabulary = {index: token for index, token in enumerate(text)}
    return vocabulary_to_index, index_to_vocabulary

test_text = '''
Moe_Szyslak Moe's Tavern Where the elite meet to drink
Bart_Simpson Eh yeah hello is Mike there Last name Rotch
Moe_Szyslak Hold on I'll check Mike Rotch Mike Rotch Hey has anybody seen Mike Rotch lately
Moe_Szyslak Listen you little puke One of these days I'm gonna catch you and I'm gonna carve my name on your back with an ice pick
Moe_Szyslak Whats the matter Homer You're not your normal effervescent self
Homer_Simpson I got my problems Moe Give me another one
Moe_Szyslak Homer hey you should not drink to forget your problems
Barney_Gumble Yeah you should only drink to enhance your social skills'''

unittests.test_create_lookup_tables(create_lookup_tables)

def token_lookup():
    """
    Generate a dict to turn punctuation into a token.

    Returns:
     Tokenized dictionary where the key is the punctuation and the value is the token
    """
    tokens = {'.': "period",
              ',': 'comma',
              '"': 'quotation',
              ';': 'semicolon',
              '!': 'exclamation',
              '?': 'question',
              '(': 'leftparenthesis',
              ')': 'rightparenthesis',
              '-': 'dash',
              '\n': 'newline'}
    return {token: '**{}**'.format(coded) for token,coded in tokens.items()}

unittests.test_tokenize(token_lookup)

text = helper.load_data(inspector.scripts.path)
text = text[81:]
token_dict = token_lookup()
for key, token in token_dict.items():
    text = text.replace(key, ' {} '.format(token))
text = text.lower()
text = text.split()
vocab_to_int, int_to_vocab = create_lookup_tables(text + list(helper.SPECIAL_WORDS.values()))
int_text = [vocab_to_int[word] for word in text]
pre_processed = inspector.scripts.path.parent.joinpath('preprocess.pkl')
with pre_processed.open("wb") as writer:
    pickle.dump((int_text, vocab_to_int, int_to_vocab, token_dict), writer)

pre_processed = inspector.scripts.path.parent.joinpath('preprocess.pkl')
with pre_processed.open("rb") as reader:
    int_text, vocab_to_int, int_to_vocab, token_dict = pickle.load(reader)

# reshape into (batch_size, seq_length, output_size)
output = output.view(batch_size, -1, self.output_size)
# get last batch
out = output[:, -1]

import torch.nn as nn

class RNN(nn.Module):

    def __init__(self, vocab_size, output_size, embedding_dim, hidden_dim, n_layers, dropout=0.5):
        """
        Initialize the PyTorch RNN Module
        :param vocab_size: The number of input dimensions of the neural network (the size of the vocabulary)
        :param output_size: The number of output dimensions of the neural network
        :param embedding_dim: The size of embeddings, should you choose to use them        
        :param hidden_dim: The size of the hidden layer outputs
        :param dropout: dropout to add in between LSTM/GRU layers
        """
        super(RNN, self).__init__()
        # TODO: Implement function

        # set class variables

        # define model layers


    def forward(self, nn_input, hidden):
        """
        Forward propagation of the neural network
        :param nn_input: The input to the neural network
        :param hidden: The hidden state        
        :return: Two Tensors, the output of the neural network and the latest hidden state
        """
        # TODO: Implement function   

        # return one batch of output word scores and the hidden state
        return None, None


    def init_hidden(self, batch_size):
        '''
        Initialize the hidden state of an LSTM/GRU
        :param batch_size: The batch_size of the hidden state
        :return: hidden state of dims (n_layers, batch_size, hidden_dim)
        '''
        # Implement function

        # initialize hidden state with zero weights, and move to GPU if available

        return None

tests.test_rnn(RNN, train_on_gpu)

loss = forward_back_prop(decoder, decoder_optimizer, criterion, inp, target)

def forward_back_prop(rnn, optimizer, criterion, inp, target, hidden):
    """
    Forward and backward propagation on the neural network
    :param decoder: The PyTorch Module that holds the neural network
    :param decoder_optimizer: The PyTorch optimizer for the neural network
    :param criterion: The PyTorch loss function
    :param inp: A batch of input to the neural network
    :param target: The target output for the batch of input
    :return: The loss and the latest hidden state Tensor
    """

    # TODO: Implement Function

    # move data to GPU, if available

    # perform backpropagation and optimization

    # return the loss over a batch and the hidden state produced by our model
    return None, None

# Note that these tests aren't completely extensive.
# they are here to act as general checks on the expected outputs of your functions
"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
tests.test_forward_back_prop(RNN, forward_back_prop, train_on_gpu)

# Data params
# Sequence Length
sequence_length =   # of words in a sequence
# Batch Size
batch_size = 

# data loader - do not change
train_loader = batch_data(int_text, sequence_length, batch_size)

# Number of Epochs
num_epochs = 
# Learning Rate
learning_rate = 

# Model parameters
# Vocab size
vocab_size = 
# Output size
output_size = 
# Embedding Dimension
embedding_dim = 
# Hidden Dimension
hidden_dim = 
# Number of RNN Layers
n_layers = 

# Show stats for every n number of batches
show_every_n_batches = 500

# create model and move to gpu if available
rnn = RNN(vocab_size, output_size, embedding_dim, hidden_dim, n_layers, dropout=0.5)
if train_on_gpu:
    rnn.cuda()

# defining loss and optimization functions for training
optimizer = torch.optim.Adam(rnn.parameters(), lr=learning_rate)
criterion = nn.CrossEntropyLoss()

# training the model
trained_rnn = train_rnn(rnn, batch_size, optimizer, criterion, num_epochs, show_every_n_batches)

# saving the trained model
helper.save_model('./save/trained_rnn', trained_rnn)
print('Model Trained and Saved')

import torch
import helper
import problem_unittests as tests

_, vocab_to_int, int_to_vocab, token_dict = helper.load_preprocess()
trained_rnn = helper.load_model('./save/trained_rnn')