# ---
# 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
# ---

# + _cell_guid="b1076dfc-b9ad-4769-8c92-a6c4dae69d19" _uuid="8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load in 

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the "../input/" directory.
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        
        print(os.path.join(dirname, filename))

# Any results you write to the current directory are saved as output.

# +
import yaml
df = pd.DataFrame(columns = ['question', 'answer']) 
idx = 0
def add_to_df(file_name):
    global idx
    with open(r'/kaggle/input/chatterbotenglish/' + file_name) as file:
        content_list = yaml.load(file, Loader=yaml.FullLoader)
        conversation = content_list['conversations']
        #print(content_list['conversations'][0])
        for i in range(len(conversation)):
            q = conversation[i][0]
            for j in range(1,len(conversation[i])):
                a = conversation[i][j]   
                df.loc[idx] = [q,a]
                idx+=1
        #print(content_list['conversations'])


for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        if (filename[-3:]=="yml"):
            if (filename != "literature.yml"):
                add_to_df(str(filename))

#add_to_df('emotion.yml')
df

# + _cell_guid="79c7e3d0-c299-4dcb-8224-4455121ee9b0" _uuid="d629ff2d2480ee46fbb7e2d37f6b5fab8052498a"
import torch
from torch.utils.data import Dataset
from transformers import BertTokenizer
import pandas as pd

class ChatbotDataset(Dataset):

    def __init__(self, df, maxlen):

        #Store the contents of the file in a pandas dataframe
        self.df = df

        #Initialize the BERT tokenizer
        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

        self.maxlen = maxlen

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

    def __getitem__(self, index):
        #print (index)
        #Selecting the question and answer at the specified index in the data frame
        question = self.df.loc[index, 'question']
        answer = self.df.loc[index, 'answer']
        #print (question)
        #print (answer)
        #Preprocessing the text to be suitable for BERT
        tokens = self.tokenizer.tokenize(question) #Tokenize the sentence
        tokens = ['[CLS]'] + tokens + ['[SEP]'] #Insering the CLS and SEP token in the beginning and end of the sentence
        if len(tokens) < self.maxlen:
            tokens = tokens + ['[PAD]' for _ in range(self.maxlen - len(tokens))] #Padding sentences
        else:
            tokens = tokens[:self.maxlen-1] + ['[SEP]'] #Prunning the list to be of specified max length
        #print (tokens)
        tokens_ids = self.tokenizer.convert_tokens_to_ids(tokens) #Obtaining the indices of the tokens in the BERT Vocabulary
        #print (tokens_ids)
        tokens_ids_tensor = torch.tensor(tokens_ids) #Converting the list to a pytorch tensor

        #Obtaining the attention mask i.e a tensor containing 1s for no padded tokens and 0s for padded ones
        attn_mask = (tokens_ids_tensor != 0).long()
        
        ans_tokens = self.tokenizer.tokenize(answer)
        #print (ans_tokens)
        ans_tokens = ['[CLS]'] + ans_tokens + ['[SEP]'] #Insering the CLS and SEP token in the beginning and end of the sentence
        ans_len = []
        if len(ans_tokens) < self.maxlen:
            ans_tokens = ans_tokens + ['[PAD]' for _ in range(self.maxlen - len(ans_tokens))] #Padding sentences
            ans_len = [len(ans_tokens)]
        else:
            ans_tokens = ans_tokens[:self.maxlen-1] + ['[SEP]'] #Prunning the list to be of specified max length
            ans_len = [len(ans_tokens)]
        ans_tokens_ids = self.tokenizer.convert_tokens_to_ids(ans_tokens) #Obtaining the indices of the tokens in the BERT Vocabulary
        #print (tokens_ids)
        ans_tokens_ids_tensor = torch.tensor(ans_tokens_ids) #Converting the list to a pytorch tensor
        
        ans_len = torch.LongTensor(ans_len)
        #print (self.tokenizer.vocab_size)
        return tokens_ids_tensor, attn_mask, ans_tokens_ids_tensor,ans_len


# +
from torch.utils.data import DataLoader

#Creating instances of training set
train_set = ChatbotDataset(df, maxlen = 15)

# +
import torch.nn as nn
from transformers import BertModel

class Encoder(nn.Module):

    def __init__(self, encode_dim, freeze_bert = True):
        super(Encoder, self).__init__()
        #Instantiating BERT model object 
        self.bert_layer = BertModel.from_pretrained('bert-base-uncased')
        self.encode_dim = encode_dim
        #Freeze bert layers
        if freeze_bert:
            for p in self.bert_layer.parameters():
                p.requires_grad = False
        
        self.fc = nn.Linear(768,encode_dim)
        

    def forward(self, seq, attn_masks):
        '''
        Inputs:
            -seq : Tensor of shape [B, T] containing token ids of sequences
            -attn_masks : Tensor of shape [B, T] containing attention masks to be used to avoid contibution of PAD tokens
        '''

        #Feeding the input to BERT model to obtain contextualized representations
        cont_reps, _ = self.bert_layer(seq, attention_mask = attn_masks)
        #print (cont_reps.size())  (batch_size , max_len, 768) 
        self.out = self.fc(cont_reps) # (batch_size, max_len, encode_dim)
        #print (self.out.size())
        return self.out
# -

net = Encoder(512,freeze_bert = True)


class Attention(nn.Module):
    """
    Attention Network.
    """

    def __init__(self, encoder_dim, decoder_dim):
        """
        :param encoder_dim: feature size of encoded images = 2048
        :param decoder_dim: size of decoder's RNN 
        """
        super(Attention, self).__init__()
        self.encoder_att = nn.Linear(encoder_dim, decoder_dim)  # linear layer to transform encoded vector
        self.softmax = nn.Softmax(dim=1)  # softmax layer to calculate weights

    def forward(self, encoder_out, decoder_hidden):
        """
        Forward propagation.
        :param encoder_out: encoded vector, a tensor of dimension (batch_size, max_len, encoder_dim)
        :param decoder_hidden: previous decoder output, a tensor of dimension (batch_size, decoder_dim)
        :return: attention weighted encoding, weights
        """
        #print (encoder_out.size())
        att1 = self.encoder_att(encoder_out)  # (batch_size, max_len, decoder_dim)
        decoder_hidden_copy=decoder_hidden.unsqueeze(2)  #(batch_size, decoder_dim,1)
        #print (encoder_out.size())
        #print (decoder_hidden_copy.size())
        score = torch.bmm(att1,decoder_hidden_copy)  # (batch_size , max_len ,1)
        
        score = score.squeeze(2) # (batch_size , max_len)
        
        alpha = self.softmax(score)  # (batch_size, max_len)
        attention_weighted_encoding = (encoder_out * alpha.unsqueeze(2)).sum(dim=1)  # (batch_size, encoder_dim)

        return attention_weighted_encoding, alpha


class DecoderWithAttention(nn.Module):
    """
    Decoder.
    """

    def __init__(self, embed_dim, decoder_dim, vocab_size, encoder_dim=256, dropout=0.5):
        """
        :param embed_dim: embedding size
        :param decoder_dim: size of decoder's RNN
        :param vocab_size: size of vocabulary
        :param encoder_dim: feature size of encoded vector
        :param dropout: dropout
        """
        super(DecoderWithAttention, self).__init__()

        self.encoder_dim = encoder_dim
        self.embed_dim = embed_dim
        self.decoder_dim = decoder_dim
        self.vocab_size = vocab_size
        self.dropout = dropout

        self.attention = Attention(encoder_dim, decoder_dim)  # attention network

        self.embedding = nn.Embedding(vocab_size, embed_dim)  # embedding layer
        self.dropout = nn.Dropout(p=self.dropout)
        self.decode_step = nn.LSTMCell(embed_dim + encoder_dim, decoder_dim, bias=True)  # decoding LSTMCell
        self.init_h = nn.Linear(encoder_dim, decoder_dim)  # linear layer to find initial hidden state of LSTMCell
        self.init_c = nn.Linear(encoder_dim, decoder_dim)  # linear layer to find initial cell state of LSTMCell
        self.f_beta = nn.Linear(decoder_dim, encoder_dim)  # linear layer to create a sigmoid-activated gate
        self.sigmoid = nn.Sigmoid()
        self.fc = nn.Linear(decoder_dim, vocab_size)  # linear layer to find scores over vocabulary
        self.init_weights()  

    def init_weights(self):
        self.embedding.weight.data.uniform_(-0.1, 0.1)
        self.fc.bias.data.fill_(0)

    def init_hidden_state(self, encoder_out):
        """
        Creates the initial hidden and cell states for the decoder's LSTM.
        :param encoder_out: encoded vector, a tensor of dimension (batch_size, max_len, encoder_dim)
        :return: hidden state, cell state
        """
        mean_encoder_out = encoder_out.mean(dim=1)
        h = self.init_h(mean_encoder_out)  # (batch_size, decoder_dim)
        c = self.init_c(mean_encoder_out)
        return h, c

    def forward(self, encoder_out, encoded_captions, caption_lengths):
        """
        Forward propagation.
        :param encoder_out: encoded vector, a tensor of dimension (batch_size, max_len, encoder_dim)
        :param encoded_captions: encoded captions, a tensor of dimension (batch_size, max_caption_length)
        :param caption_lengths: caption lengths, a tensor of dimension (batch_size, 1)
        :return: scores for vocabulary, sorted encoded captions, decode lengths, weights, sort indices
        """
        

        batch_size = encoder_out.size(0)
        vocab_size = self.vocab_size

        max_len = encoder_out.size(1)
        #caption_lengths = caption_lengths.squeeze(1)
        #print (caption_lengths.size())
        #print (encoder_out.size())
        caption_lengths, sort_ind = caption_lengths.squeeze(1).sort(dim=0, descending=True)
        #print (encoder_out.size())
        #print (sort_ind)
        
        encoder_out = encoder_out[sort_ind]
        #print (encoder_out.size())
        encoded_captions = encoded_captions.squeeze(1)
        #print (encoded_captions)
        #print (encoded_captions.size())
        
        encoded_captions = encoded_captions[sort_ind]
        #print (encoded_captions)
        #print (encoded_captions.size())
        # Embedding
        embeddings = self.embedding(encoded_captions)  # (batch_size, max_caption_length, embed_dim)
        #print (embeddings.size())
        # Initialize LSTM state
        h, c = self.init_hidden_state(encoder_out)  # (batch_size, decoder_dim)

        decode_lengths = (caption_lengths - 1).tolist()
        #print (caption_lengths)
        #print (decode_lengths)
        '''
        sample_decode_lengths=[]
        for i in decode_lengths:
            sample_decode_lengths.append(max(i))
        #print (sample_decode_lengths)
        decode_lengths = sample_decode_lengths
        '''
        # Create tensors to hold word predicion scores and alphas
        predictions = torch.zeros(batch_size, max(decode_lengths), vocab_size).to(device)
        alphas = torch.zeros(batch_size, max(decode_lengths), max_len).to(device)

        # At each time-step, decode by
        # attention-weighing the encoder's output based on the decoder's previous hidden state output
        # then generate a new word in the decoder with the previous word and the attention weighted encoding
        
        #print (max(decode_lengths))
        for t in range(max(decode_lengths)):
            batch_size_t = sum([l > t for l in decode_lengths])
            
            attention_weighted_encoding, alpha = self.attention(encoder_out[:batch_size_t],
                                                                h[:batch_size_t])
            gate = self.sigmoid(self.f_beta(h[:batch_size_t]))  # gating scalar, (batch_size_t, encoder_dim)
            #print ("BEFORE GATE")
            #print (attention_weighted_encoding.size())
            attention_weighted_encoding = gate * attention_weighted_encoding
            #print ("AFTER GATE")
            #print (embeddings.size())
            #print (attention_weighted_encoding.size())
            h, c = self.decode_step(
                torch.cat([embeddings[:batch_size_t, t, :], attention_weighted_encoding], dim=1),
                (h[:batch_size_t], c[:batch_size_t]))  # (batch_size_t, decoder_dim)
            preds = self.fc(self.dropout(h))  # (batch_size_t, vocab_size)
            predictions[:batch_size_t, t, :] = preds
            alphas[:batch_size_t, t, :] = alpha

        return predictions, encoded_captions, decode_lengths, alphas, sort_ind

# +
train_loader = DataLoader(train_set, batch_size = 16, num_workers = 5)

for epoch in range(1):
    for it, (seq, attn_masks, ans, ans_len) in enumerate(train_loader):
        logits = net(seq, attn_masks)
        print (ans.size())
        print (ans_len.size())
        break

# +
# Training

import time
import torch.backends.cudnn as cudnn
import torch.optim
import torch.nn.functional as F
import torch.utils.data
import torchvision.transforms as transforms
from torch import nn
from torch.nn.utils.rnn import pack_padded_sequence

# Model parameters
emb_dim = 256  # dimension of word embeddings
decoder_dim = 256  # dimension of decoder RNN
dropout = 0.5
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print (device)
# Training parameters
start_epoch = 0
epochs = 10  
epochs_since_improvement = 0  
batch_size = 16
workers = 1  
decoder_lr = 4e-2   
alpha_c = 1.  # regularization parameter for 'doubly stochastic attention', as in the paper
best_bleu4 = 0.  # BLEU-4 score right now
print_freq = 100  
checkpoint = '../input/chatbot/checkpoint1.pt'  #None # path to checkpoint, None if none
data_name='1'
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
#rev_word_map = {v: k for k, v in word_map.items()} # ix2word
def save_checkpoint(data_name, epoch, epochs_since_improvement, encoder, decoder, encoder_optimizer, decoder_optimizer):
    state = {'epoch': epoch,
             'epochs_since_improvement': epochs_since_improvement,
             'encoder_state_dict': encoder.state_dict(),
             'decoder_state_dict': decoder.state_dict(),
             'decoder_optimizer_dict': decoder_optimizer.state_dict(),
             'encoder_optimizer_dict': encoder_optimizer.state_dict()
            }
    filename = 'checkpoint1.pt'
    print (filename)
    torch.save(state, filename)

def adjust_learning_rate(optimizer, shrink_factor):
    """
    Shrinks learning rate by a specified factor.
    :param optimizer: optimizer whose learning rate must be shrunk.
    :param shrink_factor: factor in interval (0, 1) to multiply learning rate with.
    """
    print("\nDECAYING learning rate.")
    for param_group in optimizer.param_groups:
        param_group['lr'] = param_group['lr'] * shrink_factor
    print("The new learning rate is %f\n" % (optimizer.param_groups[0]['lr'],))    
    

def train(train_loader, encoder, decoder, criterion, encoder_optimizer,decoder_optimizer, epoch , start ):
    decoder.train()  
    encoder.train()
    
    for i, (seq, attn_masks, ans, ans_len) in enumerate(train_loader):
        current_time = time.time()
        if (current_time-start >= 6*3600):
            break
        
        # Move to GPU, if available
        seq = seq.to(device)
        attn_masks = attn_masks.to(device)
        ans = ans.to(device)
        ans_len = ans_len.to(device)
        print (seq.size())
        # Forward prop.
        out = encoder(seq, attn_masks)
        scores, caps_sorted, decode_lengths, alphas, sort_ind = decoder(out, ans, ans_len)

        # Since we decoded starting with <start>, the targets are all words after <start>, up to <end>
        #print (caps_sorted.size())
        
        targets = (caps_sorted[:, 1:].clone())
        #print (targets)
        
        if (epoch%1==0):
            with torch.no_grad():
                ans=[]
                for t in range(max(decode_lengths)):
                    #print (scores.size())
                    data = scores[0,t,:]
                    data = F.softmax(data, dim=0)
                    #print (data.size())

                    val = (max(data))
                    for j in range(30522):
                        if data[j]==val:
                            predicted_token = tokenizer.convert_ids_to_tokens([j])[0]
                            ans.append(predicted_token)
                            break
                Ans=[]
                #print (scores.view(-1,30522).size())
                #print (targets.squeeze(1).size())
                for j in (targets[0]):
                    #print (j.data.tolist())
                    Ans.append(tokenizer.convert_ids_to_tokens([j.data.tolist()]))
                print (ans)
                print (Ans)
        # Remove timesteps that we didn't decode at, or are pads
        #scores, _ = pack_padded_sequence(scores, decode_lengths, batch_first=True)
        #targets, _ = pack_padded_sequence(targets, decode_lengths, batch_first=True)
        
        
        
        #print (scores.size())
        #print (targets.size())
        t = targets
        #print (t.view(-1,1).size())
        
        # Calculate loss
        loss = criterion(scores.view(-1,30522), targets.view(-1,1).squeeze(1))

        # Add doubly stochastic attention regularization
        loss += alpha_c * ((1. - alphas.sum(dim=1)) ** 2).mean()
        print (loss.item())
        # Back prop.
        encoder_optimizer.zero_grad()
        decoder_optimizer.zero_grad()
        loss.backward()

        
        # Update weights
        encoder_optimizer.step()
        decoder_optimizer.step()
        

def main():
    global epochs_since_improvement, checkpoint, start_epoch, data_name, word_map
    
    
    decoder = DecoderWithAttention(embed_dim=emb_dim,
                                   decoder_dim=decoder_dim,
                                   vocab_size=30522,
                                   dropout=dropout)
    decoder_optimizer = torch.optim.Adam(params=filter(lambda p: p.requires_grad, decoder.parameters()),
                                         lr=decoder_lr)
    
    encoder = Encoder(256,freeze_bert = True)
    encoder_optimizer = torch.optim.Adam(params=filter(lambda p: p.requires_grad, encoder.parameters()),
                                         lr=decoder_lr)
    # Move to GPU, if available
    decoder = decoder.to(device)
    encoder = encoder.to(device)
    if checkpoint is not None:
        checkpoint = torch.load(checkpoint)
        start_epoch = checkpoint['epoch'] + 1
        epochs_since_improvement = checkpoint['epochs_since_improvement']
        decoder.load_state_dict(checkpoint['decoder_state_dict'])
        
        decoder_optimizer.load_state_dict(checkpoint['decoder_optimizer_dict'])
        encoder.load_state_dict(checkpoint['encoder_state_dict'])
        encoder_optimizer.load_state_dict(checkpoint['encoder_optimizer_dict'])

    # Loss function
    criterion = nn.CrossEntropyLoss().to(device)
    
    dataloader = DataLoader(train_set, batch_size = 16, num_workers = 0)
    start = time.time()
    while(1):
        for epoch in range(start_epoch, epochs):


            current_time = time.time()
            if (current_time-start >= 6*3600):
                print ("2 minutes completed ")
                print (epoch)
                #save_checkpoint(data_name, epoch, epochs_since_improvement, encoder, decoder, decoder_optimizer)
                break
            train(dataloader, encoder, decoder, criterion,encoder_optimizer, decoder_optimizer, epoch ,start)

        print ("EPOCHS COMPLETED")
        save_checkpoint(data_name, 0, epochs_since_improvement, encoder, decoder,encoder_optimizer, decoder_optimizer)
        if (current_time-start >= 6*3600):
            print ("2 minutes completed ")
            print (epoch)
            #save_checkpoint(data_name, epoch, epochs_since_improvement, encoder, decoder, decoder_optimizer)
            break
#main()


# +
## EVALUATING ##


decoder = DecoderWithAttention(embed_dim=emb_dim,
                                   decoder_dim=decoder_dim,
                                   vocab_size=30522,
                                   dropout=dropout)

encoder = Encoder(256,freeze_bert = True)
decoder = decoder.to(device)
encoder = encoder.to(device)
checkpoint = '../input/chatbot/checkpoint1.pt'
if checkpoint is not None:
    checkpoint = torch.load(checkpoint)
    decoder.load_state_dict(checkpoint['decoder_state_dict'])
    encoder.load_state_dict(checkpoint['encoder_state_dict'])
    

# +
encoder = encoder.eval()
decoder = decoder.eval()

question = "chemistry"
maxlen = 15
tokens = tokenizer.tokenize(question) #Tokenize the sentence
tokens = ['[CLS]'] + tokens + ['[SEP]'] #Insering the CLS and SEP token in the beginning and end of the sentence
if len(tokens) < maxlen:
    tokens = tokens + ['[PAD]' for _ in range(maxlen - len(tokens))] #Padding sentences
else:
    tokens = tokens[:maxlen-1] + ['[SEP]'] #Prunning the list to be of specified max length
#print (tokens)
tokens_ids = tokenizer.convert_tokens_to_ids(tokens) #Obtaining the indices of the tokens in the BERT Vocabulary
#print (tokens_ids)
tokens_ids_tensor = torch.tensor(tokens_ids) #Converting the list to a pytorch tensor

#Obtaining the attention mask i.e a tensor containing 1s for no padded tokens and 0s for padded ones
attn_mask = (tokens_ids_tensor != 0).long().unsqueeze(0)
seq = tokens_ids_tensor.unsqueeze(0)

# Forward prop.
out = encoder(seq, attn_mask)
pred = torch.LongTensor([[tokenizer.convert_tokens_to_ids('[CLS]')]]).to(device)
h,c = decoder.init_hidden_state(out) 

sampled = []
for t in range(20):
    embeddings = decoder.embedding(pred).squeeze(1)  
    attention_weighted_encoding, alpha = decoder.attention(out,h)
    gate = decoder.sigmoid(decoder.f_beta(h))  
    attention_weighted_encoding = gate * attention_weighted_encoding
    h, c = decoder.decode_step(
        torch.cat([embeddings, attention_weighted_encoding], dim=1),(h, c))  # (batch_size_t, decoder_dim)
    pt = decoder.fc(decoder.dropout(h))  # (batch_size_t, vocab_size)
    _,pred = pt.max(1)
    sampled.append(pred.item())

#print (sampled)
print (question)
answer = []
for i in sampled:
    answer.append(tokenizer.convert_ids_to_tokens(i))
    if (answer[-1]=='[SEP]'):
        break

print (answer)
# -

x = torch.arange(0,5)
torch.nn.functional.one_hot(x)
from torch.autograd import Variable
output = (torch.rand(14,10))
target = (torch.LongTensor(14))
for i in range(14):
    target[i] = i%10
print (output.size())
print (output.view(-1,1).size())
print (target.size())
criterion = nn.CrossEntropyLoss()
loss = criterion(output, target)
print(loss)

StarcoderdataJupyterScriptsDedupFiltered