LSTM Pytorch

from collections import defaultdict
import copy
import random
import os
import shutil
from urllib.request import urlretrieve
import numpy as np
import matplotlib.pyplot as plt
from tqdm.notebook import tqdm
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.nn.functional as F
import torch.optim
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
from multiprocessing import Pool
cudnn.benchmark = True
from torchsummary import summary
import pickle
import pandas as pd
import time
import random

#fix seed for stable training
SEED = 42

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.backends.cudnn.deterministic = True

LEANRING_RATE = 0.001
#device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
epochs = 30

#read data
df_text = pd.read_csv("grimms_fairytales.csv")
list_all_text = []
for i in range(df_text.shape[0]):
  list_all_text.append(df_text['Text'][i])
#get all stories 
text = " ".join(list_all_text)

#drop new line mask
text = text.replace("\n", "")

#int2char, which maps integers to characters ; char2int, which maps characters to unique integers
chars = tuple(set(text))
int2char = dict(enumerate(chars))
char2int = {ch: ii for ii, ch in int2char.items()}
#encode char to interge
encoded = np.array([char2int[ch] for ch in text])

# one-hot encoded for each character 
def one_hot_encode(arr, n_labels):
    
    # Initialize the the encoded array
    one_hot = np.zeros((np.multiply(*arr.shape), n_labels), dtype=np.float32)
    
    # Fill the appropriate elements with ones
    one_hot[np.arange(one_hot.shape[0]), arr.flatten()] = 1.
    
    # Finally reshape it to get back to the original array
    one_hot = one_hot.reshape((*arr.shape, n_labels))
    
    return one_hot

# get data batch
def get_batches(arr, n_seqs, n_steps):
    '''Create a generator that returns batches of size
       n_seqs x n_steps from arr.
       
       Arguments
       ---------
       arr: Array you want to make batches from
       n_seqs: Batch size, the number of sequences per batch
       n_steps: Number of sequence steps per batch
    '''
    
    batch_size = n_seqs * n_steps
    n_batches = len(arr)//batch_size
    
    # Keep only enough characters to make full batches
    arr = arr[:n_batches * batch_size]
    # Reshape into n_seqs rows
    arr = arr.reshape((n_seqs, -1))
    
    for n in range(0, arr.shape[1], n_steps):
        # The features
        x = arr[:, n:n+n_steps]
        # The targets, shifted by one
        y = np.zeros_like(x)
        try:
            y[:, :-1], y[:, -1] = x[:, 1:], arr[:, n+n_steps]
        except IndexError:
            y[:, :-1], y[:, -1] = x[:, 1:], arr[:, 0]
        yield x, y

class LSTM(nn.Module):
    
    def __init__(self, tokens, device, n_steps=100, n_hidden=256, n_layers=2, drop_prob=0.3):
        super().__init__()
        self.drop_prob = drop_prob
        self.n_layers = n_layers
        self.n_hidden = n_hidden
        self.device = device

        # creating character dictionaries
        self.chars = tokens
        self.int2char = dict(enumerate(self.chars))
        self.char2int = {ch: ii for ii, ch in self.int2char.items()}
        
        ## the LSTM
        self.lstm = nn.LSTM(len(self.chars), n_hidden, n_layers, 
                            dropout=drop_prob, batch_first=True)
        
        ##  dropout layer
        self.dropout = nn.Dropout(drop_prob)
        
        ## fully-connected output layer
        self.fc = nn.Linear(n_hidden, len(self.chars))
        
        # initialize the weights
        self.init_weights()
      
    
    def forward(self, x, hc):
        ''' Forward pass through the network. 
            These inputs are x, and the hidden/cell state `hc`. '''
        
        ## Get x, and the new hidden state (h, c) from the lstm
        x, (h, c) = self.lstm(x, hc)
        
        ## pass x through a droupout layer
        x = self.dropout(x)
        
        # Stack up LSTM outputs using view
        x = x.reshape(x.size()[0]*x.size()[1], self.n_hidden)
        
        ## put x through the fully-connected layer
        x = self.fc(x)
        
        # return x and the hidden state (h, c)
        return x, (h, c)
    
    def predict(self, char, h=None, device=device, top_k=5):
      ''' Given a character, predict the next character.
      
          Returns the predicted character and the hidden state.
      '''
      
      
      if h is None:
          h = self.init_hidden(1)
      
      x = np.array([[self.char2int[char]]])
      x = one_hot_encode(x, len(self.chars))
      inputs = torch.from_numpy(x)
      inputs = inputs.to(device)
      
      h = tuple([each.data for each in h])
      out, h = self.forward(inputs, h)

      p = F.softmax(out, dim=1).data
      p = p.to(device)
      
      if top_k is None:
          top_ch = np.arange(len(self.chars))
      else:
          p, top_ch = p.topk(top_k)
          top_ch = top_ch.cpu().numpy().squeeze()
      
      p = p.cpu().numpy().squeeze()
      char = np.random.choice(top_ch, p=p/p.sum())
          
      return self.int2char[char], h
    
    def init_weights(self):
        ''' Initialize weights for fully connected layer '''
        initrange = 0.1
        
        # Set bias tensor to all zeros
        self.fc.bias.data.fill_(0)
        # FC weights as random uniform
        self.fc.weight.data.uniform_(-1, 1)
        
    def init_hidden(self, n_seqs):
        ''' Initializes hidden state '''
        # Create two new tensors with sizes n_layers x n_seqs x n_hidden,
        # initialized to zero, for hidden state and cell state of LSTM
        weight = next(self.parameters()).data
        return (weight.new(self.n_layers, n_seqs, self.n_hidden).zero_(),
                weight.new(self.n_layers, n_seqs, self.n_hidden).zero_())

# define and print the model
model = LSTM(chars, device, n_hidden=512, n_layers=2)
print(model)

#optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=LEANRING_RATE)
#loss function
criterion = nn.CrossEntropyLoss()

def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs

def train(model, data, epochs=epochs, n_seqs=10, n_steps=50, opt=optimizer, criterion=criterion, device=device):
    ''' 
    Training a text generation network 
    '''
    list_train_loss = []
    list_val_loss = []
    model.train()
    
    # create training and validation data
    val_frac = 0.1 # ratio train/val
    print_every = 50 #logging each n steps
    clip = 5 # gradient clipping
    best_valid_loss = float('inf') # best valid loss
    val_idx = int(len(data)*(1-val_frac))
    data, val_data = data[:val_idx], data[val_idx:]
    
    model.to(device)
    
    n_chars = len(model.chars)
    for e in tqdm(range(epochs)):
        start_time = time.monotonic()
        h = model.init_hidden(n_seqs)
        for x, y in get_batches(data, n_seqs, n_steps):
        
            # One-hot encode our data and make them Torch tensors
            x = one_hot_encode(x, n_chars)
            inputs, targets = torch.from_numpy(x), torch.from_numpy(y)
            
            inputs, targets = inputs.to(device), targets.to(device)

            # Creating new variables for the hidden state, otherwise
            # we'd backprop through the entire training history
            
            h = tuple([each.data for each in h])

            model.zero_grad()
            
            output, h = model.forward(inputs, h)
            loss = criterion(output, targets.view(n_seqs*n_steps))

            loss.backward()
            
            # `clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.
            nn.utils.clip_grad_norm_(model.parameters(), clip)

            opt.step()
              
        # validation
        val_h = model.init_hidden(n_seqs)
        val_losses = []
        for x, y in get_batches(val_data, n_seqs, n_steps):
            # One-hot encode our data and make them Torch tensors
            x = one_hot_encode(x, n_chars)
            x, y = torch.from_numpy(x), torch.from_numpy(y)
            
            # Creating new variables for the hidden state, otherwise
            # we'd backprop through the entire training history
            val_h = tuple([each.data for each in val_h])
            
            inputs, targets = x, y
            inputs, targets = inputs.to(device), targets.to(device)

            output, val_h = model.forward(inputs, val_h)
            val_loss = criterion(output, targets.view(n_seqs*n_steps))
        # save training and validation loss
        val_losses.append(val_loss.item())
        list_train_loss.append(loss.item())
        list_val_loss.append(np.mean(val_losses))
        #save best model
        if np.mean(val_losses) < best_valid_loss:
            best_valid_loss = np.mean(val_losses)
            torch.save(model.state_dict(), 'best-LSTM-model.pt')
        end_time = time.monotonic()

        epoch_mins, epoch_secs = epoch_time(start_time, end_time)

        print(f'\n Epoch: {e+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')
        print(f'\tTrain Loss: {loss.item():.3f} ')
        print(f'\t Val. Loss: {np.mean(val_losses):.3f} ')
    return list_train_loss, list_val_loss

#training model
n_seqs, n_steps = 128, 100
list_train_loss, list_val_loss = train(model, encoded,  n_seqs=n_seqs, n_steps=n_steps)

# plot training & validation loss
n_epochs = range(1, epochs+1)
 
# Plot and label the training and validation loss values
# Use our custom style
plt.style.use('ggplot')
plt.plot(n_epochs, list_train_loss, 'g', label='Training Loss')
plt.plot(n_epochs, list_val_loss,  'b', label='Validation Loss')
 
# Add in a title and axes labels
plt.title('Training and Validation Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
 
# Set the tick locations
plt.xticks(range(1, epochs+1, 3))
 
# Display the plot
plt.legend(loc='best')
plt.show()

# load the best checkpoint
model.load_state_dict(torch.load('best-LSTM-model.pt'))

def text_generation(model, seq_length, input_text='a'):
    """
    Funcion to generate new character using trained deep model
    """
    model.eval()
    
    # First off, run through the prime characters
    chars = [ch for ch in input_text]
    h = model.init_hidden(1)
    for ch in input_text:
        char, h = model.predict(ch, h)

    chars.append(char)
    
    # Now pass in the previous character and get a new one
    for ii in range(seq_length):
        char, h = model.predict(chars[-1], h)
        chars.append(char)

    return ''.join(chars)

generate_text = text_generation(model, seq_length=500, input_text='story')

generate_text

class GRU(nn.Module):
    
    def __init__(self, tokens, device, n_steps=100, n_hidden=256, n_layers=2, drop_prob=0.3):
        super().__init__()
        self.drop_prob = drop_prob
        self.n_layers = n_layers
        self.n_hidden = n_hidden
        self.device = device

        # creating character dictionaries
        self.chars = tokens
        self.int2char = dict(enumerate(self.chars))
        self.char2int = {ch: ii for ii, ch in self.int2char.items()}
        
        ## the GRU
        self.gru = nn.GRU(len(self.chars), n_hidden, n_layers, 
                            dropout=drop_prob, batch_first=True)
        
        ##  dropout layer
        self.dropout = nn.Dropout(drop_prob)
        
        ## fully-connected output layer
        self.fc = nn.Linear(n_hidden, len(self.chars))
        
        # initialize the weights
        self.init_weights()
      
    
    def forward(self, x, hc):
        ''' Forward pass through the network. 
            These inputs are x, and the hidden/cell state `hc`. '''
        
        ## Get x, and the new hidden state (h, c) from the lstm
        x, h = self.gru(x, hc)
        
        ## pass x through a droupout layer
        x = self.dropout(x)
        
        # Stack up GRU outputs using view
        x = x.reshape(x.size()[0]*x.size()[1], self.n_hidden)
        
        ## put x through the fully-connected layer
        x = self.fc(x)
        
        # return x and the hidden state h
        return x, h
    
    def predict(self, char, h=None, device=device, top_k=5):
      ''' Given a character, predict the next character.
      
          Returns the predicted character and the hidden state.
      '''
      
      
      if h is None:
          h = self.init_hidden(1)
      
      x = np.array([[self.char2int[char]]])
      x = one_hot_encode(x, len(self.chars))
      inputs = torch.from_numpy(x)
      inputs = inputs.to(device)
      
      h = h.data
      h = h.to(device)
     
      out, h = self.forward(inputs, h)

      p = F.softmax(out, dim=1).data
      p = p.to(device)
      
      if top_k is None:
          top_ch = np.arange(len(self.chars))
      else:
          p, top_ch = p.topk(top_k)
          top_ch = top_ch.cpu().numpy().squeeze()
      
      p = p.cpu().numpy().squeeze()
      char = np.random.choice(top_ch, p=p/p.sum())
          
      return self.int2char[char], h
    
    def init_weights(self):
        ''' Initialize weights for fully connected layer '''
        initrange = 0.1
        
        # Set bias tensor to all zeros
        self.fc.bias.data.fill_(0)
        # FC weights as random uniform
        self.fc.weight.data.uniform_(-1, 1)
        
    def init_hidden(self, n_seqs):
        ''' Initializes hidden state '''
        # Create two new tensors with sizes n_layers x n_seqs x n_hidden,
        # initialized to zero, for hidden state and cell state of LSTM
        weight = next(self.parameters()).data
        return weight.new(self.n_layers, n_seqs, self.n_hidden).zero_()#, #weight.new(self.n_layers, n_seqs, self.n_hidden).zero_())

# define and print the model
model = GRU(chars, device, n_hidden=512, n_layers=2)
model = model.to(device)
print(model)

#optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=LEANRING_RATE)
#loss function
criterion = nn.CrossEntropyLoss()

def train_gru(model, data, epochs=epochs, n_seqs=10, n_steps=50, opt=optimizer, criterion=criterion, device=device):
    ''' 
    Training a text generation network 
    '''
    list_train_loss = []
    list_val_loss = []
    model.train()
    
    # create training and validation data
    val_frac = 0.1 # ratio train/val
    print_every = 50 #logging each n steps
    clip = 5 # gradient clipping
    best_valid_loss = float('inf') # best valid loss
    val_idx = int(len(data)*(1-val_frac))
    data, val_data = data[:val_idx], data[val_idx:]
    
    model.to(device)
    
    n_chars = len(model.chars)
    for e in tqdm(range(epochs)):
        start_time = time.monotonic()
        h = model.init_hidden(n_seqs)
      
        for x, y in get_batches(data, n_seqs, n_steps):
        
            # One-hot encode our data and make them Torch tensors
            x = one_hot_encode(x, n_chars)
            inputs, targets = torch.from_numpy(x), torch.from_numpy(y)
            
            inputs, targets = inputs.to(device), targets.to(device)

            # Creating new variables for the hidden state, otherwise
            # we'd backprop through the entire training history
            h = h.data

            model.zero_grad()
            
            output, h = model.forward(inputs, h)
            loss = criterion(output, targets.view(n_seqs*n_steps))

            loss.backward()
            
            # clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.
            nn.utils.clip_grad_norm_(model.parameters(), clip)

            opt.step()
              
        # validation
        val_h = model.init_hidden(n_seqs)
        val_losses = []
        for x, y in get_batches(val_data, n_seqs, n_steps):
            # One-hot encode our data and make them Torch tensors
            x = one_hot_encode(x, n_chars)
            x, y = torch.from_numpy(x), torch.from_numpy(y)
            
            # Creating new variables for the hidden state, otherwise
            # we'd backprop through the entire training history
            # val_h = tuple([each.data for each in val_h])
            val_h = val_h.data
            
            inputs, targets = x, y
            inputs, targets = inputs.to(device), targets.to(device)

            output, val_h = model.forward(inputs, val_h)
            val_loss = criterion(output, targets.view(n_seqs*n_steps))
        # save training and validation loss
        val_losses.append(val_loss.item())
        list_train_loss.append(loss.item())
        list_val_loss.append(np.mean(val_losses))
        #save best model
        if np.mean(val_losses) < best_valid_loss:
            best_valid_loss = np.mean(val_losses)
            torch.save(model.state_dict(), 'best-GRU-model.pt')
        end_time = time.monotonic()

        epoch_mins, epoch_secs = epoch_time(start_time, end_time)

        print(f'\n Epoch: {e+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')
        print(f'\tTrain Loss: {loss.item():.3f} ')
        print(f'\t Val. Loss: {np.mean(val_losses):.3f} ')
    return list_train_loss, list_val_loss

#training model
n_seqs, n_steps = 128, 100
list_train_loss, list_val_loss = train_gru(model, encoded, n_seqs=n_seqs, n_steps=n_steps)

# plot training & validation loss
n_epochs = range(1, epochs+1)
 
# Plot and label the training and validation loss values
# Use our custom style
plt.style.use('ggplot')
plt.plot(n_epochs, list_train_loss, 'g', label='Training Loss')
plt.plot(n_epochs, list_val_loss,  'b', label='Validation Loss')
 
# Add in a title and axes labels
plt.title('Training and Validation Loss for GRU model')
plt.xlabel('Epochs')
plt.ylabel('Loss')
 
# Set the tick locations
plt.xticks(range(1, epochs+1, 3))
 
# Display the plot
plt.legend(loc='best')
plt.show()

# load the best checkpoint
model.load_state_dict(torch.load('best-GRU-model.pt'))

generate_text = text_generation(model, seq_length=500, input_text='story')

generate_text

class RNN(nn.Module):
    
    def __init__(self, tokens, device, n_steps=100, n_hidden=256, n_layers=2, drop_prob=0.3):
        super().__init__()
        self.drop_prob = drop_prob
        self.n_layers = n_layers
        self.n_hidden = n_hidden
        self.device = device

        # creating character dictionaries
        self.chars = tokens
        self.int2char = dict(enumerate(self.chars))
        self.char2int = {ch: ii for ii, ch in self.int2char.items()}
        
        ## the RNN
        self.rnn = nn.RNN(len(self.chars), n_hidden, n_layers, 
                            dropout=drop_prob, batch_first=True)
        
        ##  dropout layer
        self.dropout = nn.Dropout(drop_prob)
        
        ## fully-connected output layer
        self.fc = nn.Linear(n_hidden, len(self.chars))
        
        # initialize the weights
        self.init_weights()
      
    
    def forward(self, x, hc):
        ''' Forward pass through the network. 
            These inputs are x, and the hidden/cell state `hc`. '''
        
        ## Get x, and the new hidden state (h, c) from the lstm
        x, h = self.rnn(x, hc)
        
        ## pass x through a droupout layer
        x = self.dropout(x)
        
        # Stack up GRU outputs using view
        x = x.reshape(x.size()[0]*x.size()[1], self.n_hidden)
        
        ## put x through the fully-connected layer
        x = self.fc(x)
        
        # return x and the hidden state h
        return x, h
    
    def predict(self, char, h=None, device=device, top_k=5):
      ''' Given a character, predict the next character.
      
          Returns the predicted character and the hidden state.
      '''
      
      
      if h is None:
          h = self.init_hidden(1)
      
      x = np.array([[self.char2int[char]]])
      x = one_hot_encode(x, len(self.chars))
      inputs = torch.from_numpy(x)
      inputs = inputs.to(device)
      
      h = h.data
      h = h.to(device)
     
      out, h = self.forward(inputs, h)

      p = F.softmax(out, dim=1).data
      p = p.to(device)
      
      if top_k is None:
          top_ch = np.arange(len(self.chars))
      else:
          p, top_ch = p.topk(top_k)
          top_ch = top_ch.cpu().numpy().squeeze()
      
      p = p.cpu().numpy().squeeze()
      char = np.random.choice(top_ch, p=p/p.sum())
          
      return self.int2char[char], h
    
    def init_weights(self):
        ''' Initialize weights for fully connected layer '''
        initrange = 0.1
        
        # Set bias tensor to all zeros
        self.fc.bias.data.fill_(0)
        # FC weights as random uniform
        self.fc.weight.data.uniform_(-1, 1)
        
    def init_hidden(self, n_seqs):
        ''' Initializes hidden state '''
        # Create two new tensors with sizes n_layers x n_seqs x n_hidden,
        # initialized to zero, for hidden state and cell state of LSTM
        weight = next(self.parameters()).data
        return weight.new(self.n_layers, n_seqs, self.n_hidden).zero_()#, #weight.new(self.n_layers, n_seqs, self.n_hidden).zero_())

# define and print the model
model = RNN(chars, device, n_hidden=512, n_layers=2)
model = model.to(device)
print(model)

def train_rnn(model, data, epochs=epochs, n_seqs=10, n_steps=50, opt=optimizer, criterion=criterion, device=device):
    ''' 
    Training a text generation network 
    '''
    list_train_loss = []
    list_val_loss = []
    model.train()
    
    # create training and validation data
    val_frac = 0.1 # ratio train/val
    print_every = 50 #logging each n steps
    clip = 5 # gradient clipping
    best_valid_loss = float('inf') # best valid loss
    val_idx = int(len(data)*(1-val_frac))
    data, val_data = data[:val_idx], data[val_idx:]
    
    model.to(device)
    
    n_chars = len(model.chars)
    for e in tqdm(range(epochs)):
        start_time = time.monotonic()
        h = model.init_hidden(n_seqs)
      
        for x, y in get_batches(data, n_seqs, n_steps):
        
            # One-hot encode our data and make them Torch tensors
            x = one_hot_encode(x, n_chars)
            inputs, targets = torch.from_numpy(x), torch.from_numpy(y)
            
            inputs, targets = inputs.to(device), targets.to(device)

            # Creating new variables for the hidden state, otherwise
            # we'd backprop through the entire training history
            h = h.data

            model.zero_grad()
            
            output, h = model.forward(inputs, h)
            loss = criterion(output, targets.view(n_seqs*n_steps))

            loss.backward()
            
            # clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.
            nn.utils.clip_grad_norm_(model.parameters(), clip)

            opt.step()
              
        # validation
        val_h = model.init_hidden(n_seqs)
        val_losses = []
        for x, y in get_batches(val_data, n_seqs, n_steps):
            # One-hot encode our data and make them Torch tensors
            x = one_hot_encode(x, n_chars)
            x, y = torch.from_numpy(x), torch.from_numpy(y)
            
            # Creating new variables for the hidden state, otherwise
            # we'd backprop through the entire training history
            # val_h = tuple([each.data for each in val_h])
            val_h = val_h.data
            
            inputs, targets = x, y
            inputs, targets = inputs.to(device), targets.to(device)

            output, val_h = model.forward(inputs, val_h)
            val_loss = criterion(output, targets.view(n_seqs*n_steps))
        # save training and validation loss
        val_losses.append(val_loss.item())
        list_train_loss.append(loss.item())
        list_val_loss.append(np.mean(val_losses))
        #save best model
        if np.mean(val_losses) < best_valid_loss:
            best_valid_loss = np.mean(val_losses)
            torch.save(model.state_dict(), 'best-RNN-model.pt')
        end_time = time.monotonic()

        epoch_mins, epoch_secs = epoch_time(start_time, end_time)

        print(f'\n Epoch: {e+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')
        print(f'\tTrain Loss: {loss.item():.3f} ')
        print(f'\t Val. Loss: {np.mean(val_losses):.3f} ')
    return list_train_loss, list_val_loss

#optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=LEANRING_RATE)
#loss function
criterion = nn.CrossEntropyLoss()

#input and ouput RNN and GRU is same, so we can utilize training function from GRU
n_seqs, n_steps = 128, 100
list_train_loss, list_val_loss = train_rnn(model, encoded, opt=optimizer, criterion=criterion, n_seqs=n_seqs, n_steps=n_steps)

# plot training & validation loss
n_epochs = range(1, epochs+1)
 
# Plot and label the training and validation loss values
# Use our custom style
plt.style.use('ggplot')
plt.plot(n_epochs, list_train_loss, 'g', label='Training Loss')
plt.plot(n_epochs, list_val_loss,  'b', label='Validation Loss')
 
# Add in a title and axes labels
plt.title('Training and Validation Loss for RNN model')
plt.xlabel('Epochs')
plt.ylabel('Loss')
 
# Set the tick locations
plt.xticks(range(1, epochs+1, 3))
 
# Display the plot
plt.legend(loc='best')
plt.show()

# load the best checkpoint
model.load_state_dict(torch.load('best-RNN-model.pt'))

generate_text = text_generation(model, seq_length=500, input_text='story')

generate_text