# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:env-TM2020] *
#     language: python
#     name: conda-env-env-TM2020-py
# ---

# +
import re
import os
import torch
import random
import torch.nn as nn

from os import listdir
from tqdm.notebook import tqdm
from collections import Counter
from gensim.models import Word2Vec
from torch.nn import functional as F
from sklearn.metrics import accuracy_score
from torch.utils.data import Dataset, DataLoader


# -

# ![Image of Yaktocat](https://i.postimg.cc/QNp5pFwd/sdf.png)
#
# The layers are as follows:
# - **Tokenize:** This is not a layer for LSTM network but a mandatory step of converting our words into tokens (integers)
# - **Embedding Layer:** that converts our word tokens (integers) into embedding of specific size
# - **LSTM Layer:** defined by hidden state dims and number of layers
# - **Fully Connected Layer:** that maps output of LSTM layer to a desired output size
# - **Sigmoid Activation Layer:** that turns all output values in a value between 0 and 1
# - **Output:** Sigmoid output from the last timestep is considered as the final output of this network

# + pycharm={"name": "#%%\n"}
class RNN(nn.Module):

    def __init__(self):
        super(RNN, self).__init__()
        
        self.embedding_dim = 64
        self.output_dim = 2

        self.hidden_dim = 25
        self.no_layers = 2

        #lstm
        self.lstm = nn.LSTM(input_size=self.embedding_dim, hidden_size=self.hidden_dim,
                           num_layers=self.no_layers, batch_first=True)

        # dropout layer
        self.dropout = nn.Dropout(0.3)

        # linear and sigmoid layer
        self.fc = nn.Linear(self.hidden_dim, self.output_dim)
        self.sig = nn.Sigmoid()

    def forward(self, x):
        batch_size = x.size(0)
        # embeddings and lstm_out

        embeds = x
        hidden = self.init_hidden(batch_size)
        hidden = tuple([each.data for each in hidden])

        #print(embeds.shape)  #[50, 500, 1000]
        lstm_out, hidden = self.lstm(embeds, hidden)

        # dropout and fully connected layer
        out = self.dropout(lstm_out)
        out = self.fc(out)

        # sigmoid function
        sig_out = self.sig(out)
        # reshape to be batch_size first

        sig_out = sig_out[:, -1] # get last batch of labels
        # return last sigmoid output and hidden state
        return sig_out

    def init_hidden(self, batch_size):
        ''' Initializes hidden state '''
        # Create two new tensors with sizes n_layers x batch_size x hidden_dim,
        # initialized to zero, for hidden state and cell state of LSTM
        h0 = torch.zeros((self.no_layers, batch_size, self.hidden_dim))
        c0 = torch.zeros((self.no_layers, batch_size, self.hidden_dim))
        hidden = (h0, c0)
        return hidden


# + [markdown] pycharm={"name": "#%% md\n"}
# Preparing dataset object and embeddings

# + pycharm={"name": "#%%\n"}
class RnnDataset(Dataset):

    def __init__(self, data, model = None):
        self.data = data
        self.length = len(self.data)
        self.words2index = {}

        padding_string = "PADDING___PADDING"
        threshold = 5
        sentences = []

        #max size sentence size for each label (total mean of the size of all the sentences of that label)
        max_size = 0
        cnt = Counter()

        for l,s in data:
            max_size += len(s)
            for w in s:
                cnt[w] += 1

        max_size = int(max_size / self.length)

        #Add padding
        new_train_data = []
        for l,s in data:
            new_s = [w if cnt[w] > threshold else "unk" for w in s]
            if len(new_s) < max_size:
                pad = [padding_string] * (max_size - len(new_s))
                new_s = new_s + pad
            elif len(new_s) > max_size:
                new_s = new_s[0:max_size]
            new_train_data.append((l,new_s))
            sentences.append(new_s)

        self.data = new_train_data

        if model is None :
            self.model = Word2Vec(sentences=sentences, min_count=-1, workers=4, size=64)
            self.model.build_vocab(cnt.keys(), update=True)
            self.model.save("w2vemb_rnn.model")
        else:
            self.model = model

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        label,sentence = self.data[index]
        sentence = [w if w in self.model else "unk" for w in sentence]
        out = torch.tensor(self.model.wv[sentence],dtype=torch.float)
        label = torch.tensor(label,dtype=torch.long)
        return out, label
    


# + [markdown] pycharm={"name": "#%% md\n"}
# preprocessing sentences: removing symbols, extra spaces, digits and tokenizing

# + pycharm={"name": "#%%\n"}
def preprocess_token(s):
    # Removing all non-word character except letters and numbers
    s = re.sub(r"[^\w\s]", '', s)

    # Replacing all extra whitespaces with no space
    s = re.sub(r"\s+", '', s)

    # replacing digits with no space
    s = re.sub(r"\d", '', s)

    return s

def tokenize(x_train):
    word_list = []

    for word in x_train.lower().split():
        word = preprocess_token(word)
        if word != '':
            word_list.append(word)

    return word_list


# + [markdown] pycharm={"name": "#%% md\n"}
# Preparing dataset instance by reading reviews and labels from files, followed by performing above defined preprocessing

# + pycharm={"name": "#%%\n"}
def get_rnn_dataset(path : str, optional_file : str = None):
    num_folder = 2
    i = 0
    data = []
    cwd = os.getcwd()
    while i < num_folder:
        label = i
        folder_name = cwd + "/" + path + "/" + str(i)
        file_names = listdir(folder_name)
        for file_name in file_names:
            file_path = folder_name + "/" + file_name
            f = open(file_path,"r",encoding="utf-8")
            s = tokenize(f.read())
            data.append((label,s))
            f.close()

        i += 1

    random.shuffle(data)

    if optional_file is None :
        return RnnDataset(data)

    return RnnDataset(data, Word2Vec.load(optional_file))


# + pycharm={"name": "#%%\n"}
def train_rnn(rnn_instance: RNN, dataloader, epochs = 10):
    # loss and optimization functions
    lr=0.005
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(rnn_instance.parameters(), lr=lr)

    clip = 5

    for epoch in range(epochs):
        iterator = tqdm(dataloader)

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

            # calculate the loss and perform backprop
            loss = criterion(output, labels)
            rnn_instance.zero_grad()
            loss.backward()

            #`clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.
            nn.utils.clip_grad_norm_(rnn_instance.parameters(), clip)
            optimizer.step()


# + pycharm={"name": "#%%\n"}
# initializing RNN instance
rnn_inst = RNN()

# reading data and preparing dataset instance
rnn_dataset = get_rnn_dataset("/train")

# + pycharm={"name": "#%%\n"}
# preparing training, validation and test set from the above obtained dataset
# train= 80% | valid = 10% | test = 10%

train_size = int(0.8 * len(rnn_dataset))
rem_size = len(rnn_dataset) - train_size
train_dataset, rem_dataset = torch.utils.data.random_split(rnn_dataset, [train_size, rem_size])

valid_size = int(rem_size/2)
test_size = rem_size - valid_size
validation_dataset, test_dataset = torch.utils.data.random_split(rem_dataset, [valid_size, test_size])

# + pycharm={"name": "#%%\n"}
# preparing dataloaders
batch_size = 50

train_loader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size)
validation_loader = DataLoader(validation_dataset, shuffle=True, batch_size=batch_size)
test_loader = DataLoader(test_dataset, shuffle=True, batch_size=batch_size)

# + pycharm={"name": "#%%\n"}
train_rnn(rnn_inst, train_loader, 10)

# + pycharm={"name": "#%%\n"}
# save model after training
torch.save(rnn_inst.state_dict(), "rnn.pt")

# + pycharm={"name": "#%%\n"}
# initializing new RNN instance using pretrained model
rnn_loaded = RNN()
rnn_loaded.load_state_dict(torch.load("rnn.pt"))
rnn_loaded.eval()


# + pycharm={"name": "#%%\n"}
# evaluation function to evaluate the model's performance w.r.t. dataset
def evaluate(clf, data_loader):

    true_labels = []
    inf_labels = []

    for data, labels in data_loader:
        out = clf(data)
        cls = torch.argmax(F.softmax(out, dim=1), dim=1)
        inf_labels.extend(cls.detach().numpy().tolist())
        true_labels.extend(labels.numpy().tolist())

    return accuracy_score(true_labels, inf_labels)


# + pycharm={"name": "#%%\n"}
evaluate(rnn_loaded, validation_loader)

# + pycharm={"name": "#%%\n"}
evaluate(rnn_loaded, test_loader)

# + pycharm={"name": "#%%\n"}


StarcoderdataJupyterScriptsDedupFiltered