# ---
# 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 torch
from torch import nn
from torch.autograd import Variable

from data_loader import DataLoader, load_dictionary
from model import UniSkip
from config import *
from datetime import datetime, timedelta
# -

d = DataLoader("./data/patent.refined.sep.unigram.id.txt")

dictionary = load_dictionary("./data/patent.refined.sep.unigram.id.txt.pkl")

len(dictionary)

spm_vocab = {}
with open('/mnt/48TB/temp/patent.refined.sep.mixed.txt.vocab') as f:
    for line in f.readlines():
        spm_vocab[line.split()[0]] = line.split()[1]
[w for w in dictionary if w not in spm_vocab]

CUDA_DEVICE

mod = UniSkip()
if USE_CUDA:
    mod.cuda(CUDA_DEVICE)

lr = 3e-4
optimizer = torch.optim.Adam(params=mod.parameters(), lr=lr)

# +
loss_trail = []
last_best_loss = None
current_time = datetime.utcnow()

def debug(i, loss, prev, nex, prev_pred, next_pred):
    global loss_trail
    global last_best_loss
    global current_time

    this_loss = loss.data[0]
    loss_trail.append(this_loss)
    loss_trail = loss_trail[-20:]
    new_current_time = datetime.utcnow()
    time_elapsed = str(new_current_time - current_time)
    current_time = new_current_time
    print("Iteration {}: time = {} last_best_loss = {}, this_loss = {}".format(
              i, time_elapsed, last_best_loss, this_loss))
    
    print("prev = {}\nnext = {}\npred_prev = {}\npred_next = {}".format(
        d.convert_var_to_sentences(prev),
        d.convert_var_to_sentences(nex),
        d.convert_var_to_sentences(prev_pred),
        d.convert_var_to_sentences(next_pred),
    ))
    
    try:
        trail_loss = sum(loss_trail)/len(loss_trail)
        if last_best_loss is None or last_best_loss > trail_loss:
            print("Loss improved from {} to {}".format(last_best_loss, trail_loss))
            
            save_loc = "./saved_models/skip-best".format(lr, VOCAB_SIZE)
            print("saving model at {}".format(save_loc))
            torch.save(mod.state_dict(), save_loc)
            
            last_best_loss = trail_loss
    except Exception as e:
       print("Couldn't save model because {}".format(e))


# +
print("Starting training...")

# a million iterations
for i in range(0, 894100):
    sentences, lengths = d.fetch_batch(32 * 10)

    loss, prev, nex, prev_pred, next_pred  = mod(sentences, lengths)
    

    if i % 100 == 0:
        debug(i, loss, prev, nex, prev_pred, next_pred)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
# -



StarcoderdataJupyterScriptsDedupFiltered