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

# +
# default_exp trainingUtils

# +
# append our Aeye package
# TODO: fix this import issue when using the package.
import sys 
sys.path.append('../aeye')

from preprocessing import Lang, tensorForImageCaption, get_preprocessed_data, SOS_token, EOS_token
from models import EncoderRNN, DecoderRNN

# +
import random
import time
import math

import torch
from torch import nn
from torch import optim

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

feature_dict, sentence_list, lang = get_preprocessed_data('train')
sentence_list = random.sample(sentence_list, len(sentence_list))

img, sent = tensorForImageCaption(feature_dict, sentence_list[3], lang)

img.shape, sent.shape

img

sent

# +
SOS_tensor = torch.tensor([[SOS_token]], device=device)
EOS_tensor = torch.tensor([[EOS_token]], device=device)

EOS_tensor.dtype

# +
teacher_forcing_ratio = 0.5
MAX_LENGTH = 30

def train(img_tensor, sent_tensor, decoder, criterion, 
          decoder_optimizer, max_length = MAX_LENGTH):
        
#     encoder_optimizer.zero_grad()
    decoder_optimizer.zero_grad()
    
    length = sent_tensor.size(0)
    img_tensor = img_tensor.unsqueeze(0)
    sent_tensor = sent_tensor.unsqueeze(1)
    loss = 0
    
    decoder_input = SOS_tensor
    decoder_hidden = img_tensor
    use_teacher_forcing = True if random.random() \
                < teacher_forcing_ratio else True
    
    if use_teacher_forcing:
        for di in range(length):
            decoder_output, decoder_hidden = decoder(
                                        decoder_input, 
                                        decoder_hidden)
            
            print('\n\n', decoder_output)
            print('\n\n', decoder_hidden)
            loss += criterion(decoder_output[0], sent_tensor[di].squeeze(0))
            decoder_input = sent_tensor[di]
            
    else:
        for di in range(max_length):
            decoder_output, decoder_hidden = decoder(
                                        decoder_input,
                                        decoder_hidden)
            
            topv, topi = decoder_output.topk(1)
            decoder_input = topi.squeeze().detach()
            loss += criterion(decoder_output[0], sent_tensor[0].squeeze(0))
    
    
    loss.backward()
    
    decoder_optimizer.step()
    
    return loss.item()/length
# -

img.dtype

# +
hidden_size = 512

decoder = DecoderRNN(hidden_size, lang.n_words)



# +
def asMinutes(s):
    m = math.floor(s / 60)
    s -= m * 60
    return '%dm %ds' % (m, s)


def timeSince(since, percent):
    now = time.time()
    s = now - since
    es = s / (percent)
    rs = es - s
    return '%s (- %s)' % (asMinutes(s), asMinutes(rs))


# +
learning_rate = 0.1

def trainIters(decoder, print_every=1000, plot_every=100,
               learning_rate=0.1):
    
    print_loss_total = 0
    start = time.time()
    
    # Get data
    feature_dict, sentence_list, lang = get_preprocessed_data('train')
    sentence_list = random.sample(sentence_list, len(sentence_list))
    n_iters = len(sentence_list)
    
    criterion = nn.NLLLoss()
    decoder_optimizer = optim.SGD(decoder.parameters(), lr = learning_rate)
    
    for iter, sentence in enumerate(sentence_list):
        img, sent = tensorForImageCaption(feature_dict, sentence, lang)
        loss = train(img, sent, decoder, criterion, decoder_optimizer)
        
        print(loss)
        break
        if (iter+1) % print_every == 0:
            print_loss_avg = print_loss_total / print_every
            print_loss_total = 0
            print('%s (%d %d%%) %.4f' % (timeSince(start, iter / n_iters),
                                        iter, iter / n_iters * 100, print_loss_avg))
            

# -

trainIters(decoder, print_every=10)



StarcoderdataJupyterScriptsDedupFiltered