character_level.py

# -*- coding: utf-8 -*-
"""character_level.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1v69aJMgxLEMksoIybnofjshIgg8DqzA3
"""

import torch
import numpy as np 
import torch.nn as nn 
import torch.nn.functional as F

#DRIVE 
import os 
from google.colab import drive
drive.mount('/content/gdrive')
os.chdir('gdrive/My Drive')

#IMPORT THE DATASET
with open('datasets/anna.txt',"r") as f:
  dataset=f.read()

dataset[:100]

#unique character vocab.
chars = tuple(set(dataset))
#maps unique charater to interger
int2char = dict(enumerate(chars))
#which maps charaters to unique intergers
char2int = {ch:ii for ii, ch in int2char.items()}

encoded = np.array([char2int[ch] for ch in dataset])

encoded[:100] #31=>C, 66=> h, 57=>a,......

def one_hot_encoder(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

test_seq=np.array([[3,5,1]])
one_hot = one_hot_encoder(test_seq,8)
print(one_hot)

def get_batches(arr,batch_size,seq_length):
  
  #number of character in batch = batch size *seq length
  n_batches = len(arr)//(batch_size*seq_length)
  #keeping only enough character to make full batches
  arr = arr[:n_batches*batch_size*seq_length]
  #reshaping the array 
  #rows= batch_size 
  #-1: placeholder 
  arr=arr.reshape((batch_size,-1))
  
  for n in range(0,arr.shape[1],seq_length):
    
    #taking all the rows, #batch_size
    #column dimension will be of seq_lenght in size
    x=arr[:,n:n+seq_length]
    
    y=np.zeros_like(x)
    
    try:
      y[:,:-1],y[:,-1] = x[:,1:], arr[:,n+seq_length]
    except IndexError:
      y[:,:-1],y[:,-1] = x[:,1:], arr[:,0]
    
    yield x,y

batches = get_batches(encoded,8,50)
x,y=next(batches)

#MODEL ARCHITECTURE
class CharacterRNN(nn.Module):
  
  def __init__(self,tokens,n_hidden=256,n_layers=2,drop_prob=.5,lr=.001):
    
    super(CharacterRNN,self).__init__()
    self.drop_prob = drop_prob
    self.n_hidden = n_hidden 
    self.n_layers= n_layers
    self.lr = lr
    
    self.chars = tokens
    self.int2char = dict(enumerate(self.chars))
    self.char2int = {ch:ii for ii, ch in self.int2char.items()}
    
    self.lstm = nn.LSTM(len(self.chars),hidden_size=n_hidden,num_layers=n_layers,dropout=drop_prob,batch_first=True)
    self.dropout=nn.Dropout(drop_prob)
    self.fc=nn.Linear(n_hidden,len(self.chars))
    
  def forward(self,x,hidden):
    
    #lstm produces lstm output and a new hidden state 
    r_output, hidden = self.lstm(x,hidden)
    
    out = self.dropout(r_output)
    
    #reshaping such that our last dimension is hidden dim
    out = out.contiguous().view(-1,self.n_hidden)
    
    out = self.fc(out)
    
    return out,hidden
  
  def init_hidden(self,batch_size):
    
    weight = next(self.parameters()).data
    
    hidden = (weight.new(self.n_layers,batch_size,self.n_hidden).zero_().cuda(),
             weight.new(self.n_layers,batch_size,self.n_hidden).zero_().cuda())
             
    return hidden

#TRAINING
def train(net,data,epochs=10,batch_size=10,seq_length=60,lr=.001,clip=5,val_frac=.1,print_every=10):
  
  net.train()
  
  optimizer = torch.optim.Adam(net.parameters(),lr=lr)
  criterion = nn.CrossEntropyLoss()
  
  val_idx = int(len(data)*(1-val_frac))
  data,val_data = data[:val_idx],data[val_idx:]
  
  net.cuda()
  counter=0
  n_chars=len(net.chars)
  for e in range(epochs):
    
    h = net.init_hidden(batch_size)
    
    for x,y in get_batches(data,batch_size,seq_length):
        counter+=1
        
        x = one_hot_encoder(x,n_chars)
        inputs , target = torch.from_numpy(x).cuda(),torch.from_numpy(y).cuda()
        
        h = tuple([each.data for each in h])
        
        net.zero_grad()
        
        output,h = net(inputs,h)
        
        loss = criterion(output,target.view(batch_size*seq_length).long())
        loss.backward()
        
        nn.utils.clip_grad_norm_(net.parameters(),clip)
        optimizer.step()
        
        if counter % print_every ==0:
          val_h = net.init_hidden(batch_size)
          val_losses = []
          net.eval()
          
          for x,y in get_batches(val_data,batch_size,seq_length):
            #print("hello")
            
            x=one_hot_encoder(x,n_chars)
            inputs, target = torch.from_numpy(x).cuda(), torch.from_numpy(y).cuda()

            val_h = tuple([each.data for each in val_h])

            #inputs, target = x,y 

            output,val_h = net(inputs,val_h)
            val_loss = criterion(output, target.view(batch_size*seq_length).long())
            #print("************",val_loss.items(),"*************")    
            val_losses.append(val_loss.item())


          net.train()
          print("Epoch: {}/{}...".format(e+1, epochs),
          "Step: {}...".format(counter),
          "Loss: {:.4f}...".format(loss.item()),
          "Val Loss: {:.4f}".format(np.mean(val_losses)))

# define and print the net
n_hidden=512
n_layers=2

net = CharacterRNN(chars, n_hidden, n_layers)
net.cuda()
print(net)

batch_size = 128 
seq_length = 100 
n_epochs = 15 

train(net,encoded,epochs=n_epochs,batch_size=batch_size,seq_length=seq_length,lr=.001,print_every=10)

os.chdir('models')

checkpoint = {
    'n_hidden':net.n_hidden,
    'n_layers':net.n_layers,
    'state_dict': net.state_dict(),
    'tokens':net.chars
}

with open("charRNN","wb") as f:
  torch.save(checkpoint,f)

with open("charRNN",'rb') as f:
  checkpoint = torch.load(f)

loaded=CharacterRNN(checkpoint['tokens'],n_hidden=checkpoint['n_hidden'],n_layers=checkpoint['n_layers'])
loaded.load_state_dict(checkpoint['state_dict'])

def predict(net,char,h=None,top_k=None):
  
  x = np.array([[net.char2int[char]]])
  x= one_hot_encoder(x,len(net.chars))
  inputs = torch.from_numpy(x).cuda()
  
  h = tuple([each.data for each in h])
  out, h = net(inputs,h)
  
  p = F.softmax(out,dim=1).data
  p=p.cpu()
  
  if top_k is None:
    top_ch = np.arange(len(net.chars))
  else:
    p, top_ch = p.topk(top_k)
    top_ch = top_ch.numpy().squeeze()
    
  p = p.numpy().squeeze()
  char = np.random.choice(top_ch,p=p/p.sum())
  
  return net.int2char[char],h

def sample(net,size,prime='The',top_k=None):
    
    net.cuda()
    net.eval()
    chars = [ch for ch in prime]
    h = net.init_hidden(1) #one character at a time hence batch size 1
    for ch in prime:
      char, h = predict(net,ch,h,top_k=top_k)
      
    chars.append(char)
    
    for ii in range(size):
      char, h = predict(net,chars[-1],h,top_k=top_k)
      chars.append(char)
     
    return ''.join(chars)

print(sample(loaded,50,prime="L",top_k=5))