<jupyter_start><jupyter_text>Base Model Seq2Seq<jupyter_code>import torch
import torch.nn as nn
import torch.optim as optim
import string
!pip install unidecode
import unidecode
import numpy as np
import itertools
from collections import Counter
from tqdm import tqdm<jupyter_output>[33mWARNING: The directory '/home/user1/.cache/pip' or its parent directory is not owned or is not writable by the current user. The cache has been disabled. Check the permissions and owner of that directory. If executing pip with sudo, you may want sudo's -H flag.[0m
Requirement already satisfied: unidecode in /home/user1/anaconda3/lib/python3.7/site-packages (1.2.0)<jupyter_text>Building blocksFirst, we will define the encode and decoder<jupyter_code>NULL_INDEX = 0

encoder_dim = 1024
predictor_dim = 1024
joiner_dim = 1024
from torch import nn
class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        
        self.hid_dim = hid_dim
        self.n_layers = n_layers
        
        self.embedding = nn.Embedding(input_dim, emb_dim)
        
        self.rnn = nn.LSTM(emb_dim, hid_dim, n_layers, dropout = dropout)
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, src):
        
        #src = [src len, batch size]
        
        embedded = self.dropout(self.embedding(src))
        
        #embedded = [src len, batch size, emb dim]
        
        outputs, (hidden, cell) = self.rnn(embedded)
        
        #outputs = [src len, batch size, hid dim * n directions]
        #hidden = [n layers * n directions, batch size, hid dim]
        #cell = [n layers * n directions, batch size, hid dim]
        
        #outputs are always from the top hidden layer
        
        return hidden, cell
class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        
        self.output_dim = output_dim
        self.hid_dim = hid_dim
        self.n_layers = n_layers
        
        self.embedding = nn.Embedding(output_dim, emb_dim)
        
        self.rnn = nn.LSTM(emb_dim, hid_dim, n_layers, dropout = dropout)
        
        self.fc_out = nn.Linear(hid_dim, output_dim)
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, input, hidden, cell):
        
        #input = [batch size]
        #hidden = [n layers * n directions, batch size, hid dim]
        #cell = [n layers * n directions, batch size, hid dim]
        
        #n directions in the decoder will both always be 1, therefore:
        #hidden = [n layers, batch size, hid dim]
        #context = [n layers, batch size, hid dim]
        
        input = input.unsqueeze(0)
        
        #input = [1, batch size]
        
        embedded = self.dropout(self.embedding(input))
        
        #embedded = [1, batch size, emb dim]
                
        output, (hidden, cell) = self.rnn(embedded, (hidden, cell))
        
        #output = [seq len, batch size, hid dim * n directions]
        #hidden = [n layers * n directions, batch size, hid dim]
        #cell = [n layers * n directions, batch size, hid dim]
        
        #seq len and n directions will always be 1 in the decoder, therefore:
        #output = [1, batch size, hid dim]
        #hidden = [n layers, batch size, hid dim]
        #cell = [n layers, batch size, hid dim]
        
        prediction = self.fc_out(output.squeeze(0))
        
        #prediction = [batch size, output dim]
        
        return prediction, hidden, cell
        
class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, device):
        super().__init__()
        
        self.encoder = encoder
        self.decoder = decoder
        self.device = device
        
        assert encoder.hid_dim == decoder.hid_dim, \
            "Hidden dimensions of encoder and decoder must be equal!"
        assert encoder.n_layers == decoder.n_layers, \
            "Encoder and decoder must have equal number of layers!"
        
    def forward(self, src, trg, teacher_forcing_ratio = 0.90):
        
        #src = [src len, batch size]
        #trg = [trg len, batch size]
        #teacher_forcing_ratio is probability to use teacher forcing
        #e.g. if teacher_forcing_ratio is 0.75 we use ground-truth inputs 75% of the time
        
        batch_size = trg.shape[1]
        trg_len = trg.shape[0]
        trg_vocab_size = self.decoder.output_dim
        
        #tensor to store decoder outputs
        outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(self.device)
        
        #last hidden state of the encoder is used as the initial hidden state of the decoder
        hidden, cell = self.encoder(src)
        
        #first input to the decoder is the <sos>
        input = trg[0,:]
        
        for t in range(1, trg_len):
            
            #insert input token embedding, previous hidden and previous cell states
            #receive output tensor (predictions) and new hidden and cell states
            output, hidden, cell = self.decoder(input, hidden, cell)
            
            #place predictions in a tensor holding predictions for each token
            outputs[t] = output
            
            #decide if we are going to use teacher forcing or not
            teacher_force = random.random() < teacher_forcing_ratio
            
            #get the highest predicted token from our predictions
            top1 = output.argmax(1) 
            
            #if teacher forcing, use actual next token as next input
            #if not, use predicted token
            input = trg[t] if teacher_force else top1
        
        return outputs
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
INPUT_DIM = 32
OUTPUT_DIM = 32
ENC_EMB_DIM = 32
DEC_EMB_DIM = 32
HID_DIM = 1024
N_LAYERS = 2
ENC_DROPOUT = 0.2
DEC_DROPOUT = 0.2
START_INDEX = 0
BLANK_INDEX = 27
PAD_INDEX = 30
END_INDEX = 31

enc = Encoder(INPUT_DIM, ENC_EMB_DIM, HID_DIM, N_LAYERS, ENC_DROPOUT)
dec = Decoder(OUTPUT_DIM, DEC_EMB_DIM, HID_DIM, N_LAYERS, DEC_DROPOUT)

model = Seq2Seq(enc, dec, device).to(device)<jupyter_output><empty_output><jupyter_text>Some utilitiesHere we will add a bit of boilerplate code for training and loading data.<jupyter_code>import random
char_labels = ['<start>', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's',
              't', 'u', 'v', 'w', 'x', 'y', 'z', '<blank>', ' ', "'", '<pad>','<end>']
char_dict = {'<start>':0, 'a':1, 'b':2, 'c':3, 'd':4, 'e':5, 'f':6, 'g':7, 'h':8, 'i':9, 'j':10, 'k':11, 'l':12, 'm':13, 
             'n':14, 'o':15, 'p':16, 'q':17, 'r':18, 's':19, 't':20, 'u':21, 'v':22, 'w':23, 'x':24, 'y':25, 'z':26, 
             '<blank>':27, ' ':28, "'":29, '<pad>':30,'<end>':31}
char_inv_dict = {0:'', 1:'a', 2:'b', 3:'c', 4:'d', 5:'e', 6:'f', 7:'g', 8:'h', 9:'i', 10:'j', 11:'k', 12:'l',
                 13:'m', 14:'n', 15:'o', 16:'p', 17:'q', 18:'r', 19:'s', 20:'t', 21:'u', 22:'v', 23:'w', 24:'x', 25:'y',
                 26:'z', 27:'', 28:' ', 29:"'", 30:'', 31:''}

class TextDataset(torch.utils.data.Dataset):
  def __init__(self, lines, batch_size):
    lines = list(filter(("\n").__ne__, lines))

    self.lines = lines # list of strings
    collate = Collate()
    self.loader = torch.utils.data.DataLoader(self, batch_size=batch_size, num_workers=0, shuffle=True, collate_fn=collate)

  def __len__(self):
    return len(self.lines)

  def __getitem__(self, idx):
    line = self.lines[idx].replace("\n", "")
    line =line.lower()
    line = "".join(c for c in line if c in char_labels)
    line = unidecode.unidecode(line) # remove special characters
    x = "".join(c*random.randint(0,1) if c in ("gjp") else c*random.randint(1,2) for c in line) 
    x = ' '.join(x.split())
    y = line
    return (x,y)

def encode_string(s):

  return [char_dict[c] for c in s]

def decode_labels(l):
  return "".join([char_inv_dict[c] for c in l])

class Collate:
  def __call__(self, batch):
    """
    batch: list of tuples (input string, output string)
    Returns a minibatch of strings, encoded as labels and padded to have the same length.
    """
    x = []; y = []
    batch_size = len(batch)
    for index in range(batch_size):
      x_,y_ = batch[index]
      x.append([0]+encode_string(x_))
      y.append([0]+encode_string(y_))

    # pad all sequences to have same length
    T = [len(x_) for x_ in x]
    U = [len(y_) for y_ in y]
 
    T_max = max(T)
    U_max = max(U)
    for index in range(batch_size):
      x[index] += [PAD_INDEX] * (T_max - len(x[index])) +[31]
      x[index] = torch.tensor(x[index])
      y[index] += [PAD_INDEX] * (U_max - len(y[index]))+[31]
      y[index] = torch.tensor(y[index])

    # stack into single tensor
    x = torch.stack(x)
    y = torch.stack(y)
    #T = torch.tensor(T)
    #U = torch.tensor(U)

    return (x,y) #,T,U)

def gen_features(filename, max_len):
    with open(filename, "r") as f:
        t_lines = f.readlines()
        t_lines = t_lines[1:]
        t_lines = [line.split('\t')[2] for line in t_lines ]
        lines = []

        for line in t_lines:
            temp = line.split(' ')
            count = len(temp)//max_len
            resid = len(temp)%max_len
    
            k = 0
            while(count > 0):    
                short_line = ' '.join(temp[k:k+max_len]) 
                lines.append(short_line)
                count = count - 1
                k = k + max_len
            if(resid > 0):
                lines.append(' '.join(temp[k:k+resid]))   
    return(lines)            


train_filename = "../LibriSpeech/train-clean-100/transcripts_460.tsv"
eval_filename = "../LibriSpeech/dev-clean/transcripts.tsv"
test_filename = "../LibriSpeech/test-clean/transcripts.tsv"


train_lines = gen_features(train_filename, 10)
eval_lines = gen_features(eval_filename, 10)
test_lines = gen_features(test_filename, 10)
train_set = TextDataset(train_lines, batch_size=32)
eval_set = TextDataset(test_lines, batch_size=8)
test_set = TextDataset(test_lines, batch_size=8)
train_set.__getitem__(0)
def init_weights(m):
    for name, param in m.named_parameters():
        nn.init.uniform_(param.data, -0.08, 0.08)
        
model.apply(init_weights)
def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f'The model has {count_parameters(model):,} trainable parameters')

optimizer = optim.Adam(model.parameters(), lr=0.0001)
criterion = nn.CrossEntropyLoss(ignore_index = PAD_INDEX)

lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=optimizer, gamma=0.98)
def train(model, iterator, optimizer, criterion, clip):
    
    model.train()
    
    epoch_loss = 0
    
    for i, batch in enumerate(iterator):
        
        src, trg = batch
        trg = trg.transpose(1,0)
        src = src.transpose(1,0)
        optimizer.zero_grad()
        src = src.to(device)
        trg = trg.to(device)
        output = model(src, trg)
        
        #trg = [trg len, batch size]
        #output = [trg len, batch size, output dim]
        
        output_dim = output.shape[-1]
        
        output = output[1:].view(-1, output_dim)
        trg = trg[1:].reshape(-1)
        
        #trg = [(trg len - 1) * batch size]
        #output = [(trg len - 1) * batch size, output dim]
        
        loss = criterion(output, trg)
        
        loss.backward()
        
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        
        optimizer.step()
        epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)



def evaluate(model, iterator, criterion):
    
    model.eval()
    
    epoch_loss = 0
    
    with torch.no_grad():
    
        for i, batch in enumerate(iterator):

            src, trg = batch
            trg = trg.transpose(1,0)
            src = src.transpose(1,0)
            src = src.to(device)
            trg = trg.to(device)
            output = model(src, trg, 0) #turn off teacher forcing

            #trg = [trg len, batch size]
            #output = [trg len, batch size, output dim]

            output_dim = output.shape[-1]
            
            output = output[1:].view(-1, output_dim)
            trg = trg[1:].reshape(-1)

            #trg = [(trg len - 1) * batch size]
            #output = [(trg len - 1) * batch size, output dim]

            loss = criterion(output, trg)
            
            epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)

N_EPOCHS = 8
CLIP = 1
best_valid_loss = 6.0 #float('inf')
train_loss = []
valid_loss = []
load_file = './base_model_2.pt'
model.load_state_dict(torch.load(load_file))
for epoch in range(N_EPOCHS):
    
    tr_loss = train(model, train_set.loader, optimizer, criterion, CLIP)
    train_loss.append(tr_loss)
    v_loss = evaluate(model, eval_set.loader, criterion)
    valid_loss.append(v_loss)
    lr_scheduler.step()
    print(tr_loss, v_loss)
    if(best_valid_loss > v_loss):
        best_valid_loss = v_loss
        save_file = './base_model.pt'
        torch.save(model.state_dict(), save_file)
print(train_loss, valid_loss)
def inference(model, iterator):
    
    model.eval()
    
    epoch_loss = 0
    y_pred = []
    y_test = []
    with torch.no_grad():
        
        for i, batch in enumerate(iterator):

            src, trg = batch
            trg = trg.transpose(1,0)
            src = src.transpose(1,0)
            src = src.to(device)
            trg = trg.to(device)
            output = model(src, trg, 0) #turn off teacher forcing

            output = output.argmax(-1)
            output = output.transpose(1,0)
            trg = trg.transpose(1,0)
            y_p = output.cpu().detach().numpy()
            y_t = trg.cpu().detach().numpy()
            y_p_s = []
            y_t_s = []
            for j in range(y_p.shape[0]):
                y_p_s.append(decode_labels(y_p[j,:]))
                y_t_s.append(decode_labels(y_t[j,:]))
            y_pred.extend(y_p_s)
            y_test.extend(y_t_s)
           
    return y_pred, y_test
y_pred, y_test = inference(model, test_set.loader)
import jiwer
from jiwer import wer
error = wer(y_test, y_pred)
print(error)<jupyter_output>0.0979664821472731
StarcoderdataJupyterStructuredCleanDedup