# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import pandas as pd
from gensim import models
import numpy as np
import gensim
import gensim.downloader
from tqdm import tqdm
import time
import math
import os

# ## Dataset

df = pd.read_csv('data/qqp/train_v1.csv', header=0, index_col=0)
df

df['question'] = (df['question1'] + "|||" + df['question2'])

df[df.is_duplicate == 1].shape[0]

df[df.is_duplicate == 0].shape[0]

# +
question_list = df['question'].to_list()

question_list = [question.lower().split() for question in question_list]

labels = df['is_duplicate'].to_list()

# +
max_length = -np.inf
for sentence in question_list:
    if len(sentence) > max_length:
        max_length = len(sentence)
        
for sentence in question_list:
    if len(sentence) > max_length:
        max_length = len(sentence)
# -

max_length

# ## Word Embeddings

root_dir = 'logs/qqp'
os.makedirs(root_dir, exist_ok=True)

# +
vector_size = 50
window_size = 5
negative_size = 15

wv_model_file = root_dir + '/' + 'wv_bilstm.pth'


# +
# wv_model = gensim.downloader.load('glove-wiki-gigaword-50')
# wv_model = models.Word2Vec(sentences=sentences, vector_size=vector_size, window=window_size, negative=negative_size).wv
# wv_model = models.Word2Vec(corpus_file='data/corpus.txt', vector_size=vector_size, window=window_size, negative=negative_size).wv

# wv_model.save(wv_model_file)
# del wv_model
# -

def vectorize_sentences(sentences, wv, sentence_size):
    voc = wv.key_to_index.keys()
    unk = wv['<unk>']
    eos = wv['<eos>']
    lengths = []
    for i, sentence in enumerate(sentences):
        lengths.append(len(sentence))
        
        for i, token in enumerate(sentence):
            if token in voc:
                sentence[i] = wv[token]
            else:
                sentence[i] = unk
        
        
        while len(sentence) < sentence_size:
            sentence.append(eos)
        
        
    return sentences, lengths

wv = models.KeyedVectors.load(wv_model_file)

wv.add_vectors(
    ['<unk>', '<eos>'],
    [np.zeros(wv.vector_size), np.ones(wv.vector_size)]
)

# ## BiLSTM

import torch
from torch import nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')


def save_model(model, file_name):
    torch.save(model.state_dict(), file_name)
def load_model(model, file_name):
    return model.load_state_dict(torch.load(file_name))


# +
hidden_size = 50
sentence_size = 275
num_layers = 2
bidirectional = True

batch_size = 128

lr = 0.001
num_epochs = 20


eval_rate = 0.1

model_file = 'logs/qqp/lstm_qqp_v2.pth'
# -

question_vec, question_lengths = vectorize_sentences(question_list, wv, sentence_size)

question_lengths = np.array(question_lengths)
labels = np.array(labels)

# +
# vec_labels = np.zeros((labels.size, labels.max()+1))
# vec_labels[np.arange(labels.size), labels] = 1

# +
# vectors = vec_sentences.reshape((-1, 50))

# mu = vectors.mean(axis=0)
# sigma = np.sqrt(((vectors - mu) ** 2).mean(axis=0))

# vec_sentences = (vec_sentences - mu) / sigma

# +
eval_index = int(len(question_vec) * eval_rate)

question_train = question_vec[eval_index:]
question_eval = question_vec[:eval_index]

question_len_train = question_lengths[eval_index:]
question_len_eval = question_lengths[:eval_index]


label_train = labels[eval_index:]
label_eval = labels[:eval_index]

# +
# question1_train = torch.tensor(question1_train, dtype=torch.float)
# question1_eval = torch.tensor(question1_eval, dtype=torch.float)

# question2_train = torch.tensor(question2_train, dtype=torch.float)
# question2_eval = torch.tensor(question2_eval, dtype=torch.float)

question_len_train = torch.tensor(question_len_train, dtype=torch.int)
question_len_eval = torch.tensor(question_len_eval, dtype=torch.int)

label_train = torch.tensor(label_train, dtype=torch.long)
label_eval = torch.tensor(label_eval, dtype=torch.long)


# -

class Classifier(nn.Module):
    
    def __init__(self):
        super(Classifier, self).__init__()
        lstm_dim = hidden_size * 2 * (2 if bidirectional else 1)
        
        
        self.lstm = nn.LSTM(input_size=vector_size,
                              hidden_size=hidden_size,
                              num_layers=num_layers,
                              bidirectional=bidirectional,
                              batch_first=True
                             )

        
        self.fcnn_1 = nn.Linear(in_features=lstm_dim, out_features=64)
        
        self.fcnn_2 = nn.Linear(in_features=64, out_features=2)
        
    def forward(self, question, question_len):
        
        question = pack_padded_sequence(question, question_len.cpu(), batch_first=True, enforce_sorted=False)
        question, _ = self.lstm(question)
        question, _ = pad_packed_sequence(question, batch_first=True)
        
        
        avg_pool = torch.mean(question, 1)
        max_pool, _ = torch.max(question, 1)
        output = torch.cat([avg_pool, max_pool], dim=1)
        
        
        output = self.fcnn_1(output)
        output = torch.relu(output)
        
        output = self.fcnn_2(output)
        return output

# +
# weights = torch.log(1/(train_y.sum(dim=0) / train_y.sum()))
# weights = weights.detach().to(device)
# weights
# -

num_train_batch = int(len(question_train) / batch_size)
num_eval_batch = int(len(question_eval) / batch_size)

classifier = Classifier().to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(classifier.parameters(), lr=lr)

# + code_folding=[]
min_loss = np.inf

for i in range(num_epochs):
    print(f'---> Epoch {i} <---')
    time.sleep(0.5)
    
    classifier.train()
    loader = tqdm(range(num_train_batch), postfix={'Epoch': i})
    train_losses = []
    
    for i_batch in loader:
        question, question_len, targets = (
            question_train[i_batch*batch_size:(i_batch+1)*batch_size],
            question_len_train[i_batch*batch_size:(i_batch+1)*batch_size],
            label_train[i_batch*batch_size:(i_batch+1)*batch_size]
        )
        
        question = torch.tensor(np.array(question), dtype=torch.float, device=device)
        
        
        targets = targets.to(device)
        
        optimizer.zero_grad()
        
        outputs = classifier(question, question_len)
        
        loss = criterion(outputs, targets)
        
        loss.backward()
        optimizer.step()
        
        with torch.no_grad():
            train_losses.append(loss.item())
            loader.set_postfix({
                'Epoch': i,
                'Train loss': np.mean(train_losses)
            }, refresh=True)
    
    
    time.sleep(0.5)
    
    with torch.no_grad():
        classifier.eval()
        loader = tqdm(range(num_eval_batch), postfix={'Epoch': i,}, colour='green')
        eval_losses = []
        eval_scores = []

        for i_batch in loader:
            question, question_len, targets = (
                question_eval[i_batch*batch_size:(i_batch+1)*batch_size],
                question_len_eval[i_batch*batch_size:(i_batch+1)*batch_size],
                label_eval[i_batch*batch_size:(i_batch+1)*batch_size]
            )
            
            question = torch.tensor(question, dtype=torch.float, device=device)
            targets = targets.to(device)
            

            outputs = classifier(question, question_len)
            
            loss = criterion(outputs, targets)
            
            score = (outputs.argmax(dim=1) == targets).detach().cpu().numpy()
            eval_scores.append(score)
            eval_losses.append(loss.item())
            loader.set_postfix({
                'Epoch': i,
                'Eval loss': np.mean(eval_losses),
                'Eval score': np.concatenate(eval_scores).mean()
            }, refresh=True)


        eval_loss = np.mean(eval_losses)
        if eval_loss <= min_loss:
            min_loss = eval_loss
            save_model(classifier, model_file)
            loader.write('*** save ***')
        
    time.sleep(0.5)
# -



StarcoderdataJupyterScriptsDedupFiltered