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

# + id="BB8HXRc_chEG" colab_type="code" colab={}
config = {
    'd_model': 200,
    'aspects_num': 50,
    'nhead': 4, # 8,
    'num_encoder_layers': 3, # 6
    'num_decoder_layers': 3, # 6
    'dim_feedforward': 2048,
    'pos_dropout': 0.2,
    'trans_dropout': 0.2,
    'batch_size': 64,
}

import os
root = "drive/My Drive/Colab Notebooks/huawei/project/generations/"
texts_path = os.path.join(root, "stop_lem_norm_sents_ria.json.gz")

# + id="WRpxX6rCdGVM" colab_type="code" outputId="3fb08b18-13ef-476c-84b9-634d040a7b7d" colab={"base_uri": "https://localhost:8080/", "height": 128}
from google.colab import drive
drive.mount('/content/drive')

# + id="-KEJUKEFM8OR" colab_type="code" outputId="890bc540-619c-41ef-d616-71e7c5cd4611" colab={"base_uri": "https://localhost:8080/", "height": 310}
# !pip install pymorphy2
# !pip install einops

# + id="okcgOCw7M-tL" colab_type="code" outputId="f1580176-d4c3-4774-f78d-b49fbb406086" colab={"base_uri": "https://localhost:8080/", "height": 164}
# %matplotlib inline
import seaborn as sns
import matplotlib.pyplot as plt


from __future__ import unicode_literals, print_function, division
from io import open
import unicodedata
import string
import re
import random
from collections import defaultdict
import numpy as np
from einops import rearrange

import nltk
from nltk.corpus import stopwords
nltk.download('stopwords')
nltk.download('punkt')

from sklearn.datasets import fetch_20newsgroups

from nltk import word_tokenize
from nltk.tokenize import sent_tokenize, word_tokenize

import pymorphy2
import re


import os
import gzip
import time
import json
import tqdm
from bs4 import BeautifulSoup

import torch
import torch.nn as nn
from torch.optim import Adam
import torch.nn.functional as F

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using {device}")


def save_texts(texts, titles, filename):
  dump_texts = os.path.join(root, filename)
  with open(file=dump_texts, mode='w') as f:
    json.dump({'texts': texts, 'titles': titles}, f)

def load_texts(filename):
  dump_texts = os.path.join(root, filename)
  with open(file=dump_texts, mode='r') as f:
    data = json.load(f)
  return data['texts'], data['titles']


# + [markdown] id="3qnoS4q_NCbT" colab_type="text"
# # Preprocess ria news

# + id="eHAyJk9sEAP1" colab_type="code" outputId="2f9d3fb3-abb4-47f4-80c0-cd270b1cb838" colab={"base_uri": "https://localhost:8080/", "height": 35}
# new
lines = []
with gzip.open(texts_path, 'rb') as f:
  for line in tqdm.tqdm(f):
    lines.append(json.loads(line))

# + id="w_5uoUEhEYeQ" colab_type="code" outputId="21e2e667-17ca-4788-f17a-5c23d4504705" colab={"base_uri": "https://localhost:8080/", "height": 53}
print(len(lines))
lines = list(filter(lambda x: len(x['title'])!=0 and len(x['title'][0])>0 and len(x['text'])>1 and len(x['text'][1])>0, lines))
print(len(lines))

# + id="mvZrDSO-EZLO" colab_type="code" outputId="305b443f-0281-4e39-989b-952b9b75063c" colab={"base_uri": "https://localhost:8080/", "height": 513}
lens = [len(l['text'][1]) for l in lines]
sns.distplot(lens)
plt.show()
plt.close()

lens = [len(l['title'][0]) for l in lines]
sns.distplot(lens)
plt.show()
plt.close()

# + id="V5vaOk64EaAA" colab_type="code" outputId="4634017e-4f1c-47b7-bab4-c1c3381b54f3" colab={"base_uri": "https://localhost:8080/", "height": 53}
print(len(lines))
lines = list(filter(lambda x: len(x['text'])!=0 and len(x['title'][0])<70 and len(x['title'][0])>30 and len(x['text'])>1 and len(x['text'][1])>5 and len(x['text'][1])<400, lines))
print(len(lines))

# + id="_iM6ZD-gEcPv" colab_type="code" outputId="28c0b21e-8497-4f9f-bfcf-ecf4fa5c3969" colab={"base_uri": "https://localhost:8080/", "height": 513}
lens = [len(l['text'][1]) for l in lines]
sns.distplot(lens)
plt.show()
plt.close()

lens = [len(l['title'][0]) for l in lines]
sns.distplot(lens)
plt.show()
plt.close()

# + id="j6txwOW8EdiF" colab_type="code" outputId="0c4bf41e-4eb4-4217-8ea2-56f98dd301e7" colab={"base_uri": "https://localhost:8080/", "height": 53}
print(len(lines))
lines = lines[:100000]
print(len(lines))

# + id="co_qy_rlEgLN" colab_type="code" outputId="318c390c-4d62-4889-bab0-f429ff7e02d1" colab={"base_uri": "https://localhost:8080/", "height": 513}
lens = [len(l['text'][1]) for l in lines]
sns.distplot(lens)
plt.show()
plt.close()

lens = [len(l['title'][0]) for l in lines]
sns.distplot(lens)
plt.show()
plt.close()

# + id="CrVZU4kvNF9x" colab_type="code" outputId="1d1a1175-6eb6-48c1-9f49-b2d8ce5e5d9e" colab={"base_uri": "https://localhost:8080/", "height": 35}
morph = pymorphy2.MorphAnalyzer()
pattern = re.compile("^[а-яА-Я]+$")

texts = []
titles = []

def process_sentence(sentence):
    words = [w for w in nltk.word_tokenize(sentence) if w not in stopwords.words("russian") and pattern.match(w)]
    words = [morph.parse(word)[0].normal_form.lower() for word in words]
    return words

for line_i in tqdm.tqdm(range(len(lines))):
    text  = lines[line_i]['text'][1] # [0] - is a paper name and date
    title = lines[line_i]['title'][0]
    processed_text = process_sentence(text)
    processed_title = process_sentence(title)
    texts.append(processed_text)
    titles.append(processed_title)

# + id="9St9F6kvHq4A" colab_type="code" colab={}
save_texts(texts, titles, "ria_text_titles.json")

# + [markdown] id="pmpGQn5VNL_p" colab_type="text"
# # Load data

# + id="B-uC9447NOjo" colab_type="code" outputId="c0c81d6c-2883-4e2e-cd4a-6bbec720c8c7" colab={"base_uri": "https://localhost:8080/", "height": 430}
texts, titles = load_texts("ria_text_titles.json") #
# texts, titles = texts[:30000], titles[:30000]
print(max([len(t) for t in texts]))



SOS_token = 0 # start of string.
EOS_token = 1 # end of string.
PAD_token = 2
UNK_token = 3

class Lang: # like a vocabulary
    def __init__(self, name):
        self.name = name
        self.word2index = {"SOS": SOS_token, "EOS": EOS_token, "PAD": PAD_token, 'UNK': UNK_token}
        self.word2count = {}
        self.index2word = {SOS_token: "SOS", EOS_token: "EOS", PAD_token: "PAD", UNK_token: "UNK"}
        self.n_words = len(self.word2index)
        self.max_seq_length = None

    def addSentence(self, sentence):
        if self.max_seq_length is None:
          self.max_seq_length = len(sentence)
        elif self.max_seq_length < len(sentence):
          self.max_seq_length = len(sentence)
        for word in sentence:
            self.addWord(word)

    def addWord(self, word):
        if word not in self.word2index:
            self.word2index[word] = self.n_words
            self.word2count[word] = 1
            self.index2word[self.n_words] = word
            self.n_words += 1
        else:
            self.word2count[word] += 1

    def tensor2text(self, tensor):
        return [self.index2word[token.item()] for token in tensor]
        
    def ids2text(self, numbers):
      return [self.index2word[num] for num in numbers]

    def text2ids(self, tokens):
        return [self.word2index[token] if token in self.word2index else self.word2index['UNK'] for token in tokens]
            
def unicodeToAscii(s):
    return ''.join(
        c for c in unicodedata.normalize('NFD', s)
        if unicodedata.category(c) != 'Mn'
    )


def normalizeString(s):
    s = unicodeToAscii(s.lower().strip())
    s = re.sub(r"([.!?])", r" \1", s)
    s = re.sub(r"[^a-zA-Z.!?]+", r" ", s)
    return s



def prepareData(texts, titles):
    pairs = [[txt, ttl] for txt, ttl in zip(texts, titles) if len(txt)>0 and len(ttl)>0 ]
    lang = Lang('ru')

    print("Read %s sentence pairs" % len(pairs))
    for pair in pairs:
        lang.addSentence(pair[0])
        lang.addSentence(pair[1])
    print("Counted words:")
    print(lang.name, lang.n_words)
    return lang, pairs

def indexesFromSentence(lang, sentence):
    return [lang.word2index[word] for word in sentence]


def tensorFromSentence(lang, sentence):
    indexes = indexesFromSentence(lang, sentence)
    indexes.append(EOS_token)
    return torch.tensor(indexes, dtype=torch.long, device=device)


def tensorsFromPair(pair):
    input_tensor = tensorFromSentence(lang, pair[0])
    target_tensor_t = tensorFromSentence(lang, pair[1][:-1])
    target_tensor_l = tensorFromSentence(lang, pair[1])
    return (input_tensor, target_tensor_t, target_tensor_l)


lang, pairs = prepareData(texts, titles)
print(random.choice(pairs))

lens = [len(t) for t in texts]
sns.distplot(lens)

i = 100
print(texts[i])
print(titles[i])
print(lang.max_seq_length)

# + [markdown] id="R0FOQBL0NU0Z" colab_type="text"
# # Datasets

# + id="Br5sXOantT0z" colab_type="code" colab={}
from torch.utils.data.dataset import Dataset
from torch.utils.data import DataLoader


class HeadlinesDataset(Dataset):
    def __init__(self, pairs, lang, batch_size):
        self.pairs = pairs
        self.lang = lang
        self.batch_size = batch_size
        self.batch_num = len(self.pairs)//batch_size
        
    def __getitem__(self, index): # todo: make number of pads depend on current batch max text, title length
        srcs = []
        srcs_lns = []
        tgts = []
        tgts_lns = []
        for i in range(self.batch_size):
          src, tgt = self.pairs[index*self.batch_size+i]
          srcs.append(src)
          srcs_lns.append(len(src))
          tgts.append(tgt)
          tgts_lns.append(len(tgt))

        max_src_ln = max(srcs_lns)          
        max_tgt_ln = max(tgts_lns)    

        converted_srcs = []    
        converted_tgts = []    
        srcs_padding_masks = []
        tgts_padding_masks = []

        for i in range(len(srcs)):
          converted_srcs.append(self.lang.text2ids(srcs[i]) + [self.lang.word2index['PAD']]*(max_src_ln-len(srcs[i]))) 
          srcs_padding_masks.append([False] + [False]*len(srcs[i]) + [True]*(max_src_ln-len(srcs[i]))) # first for intent

          converted_tgts.append(
              [self.lang.word2index['SOS']] + 
                self.lang.text2ids(tgts[i]) + 
              [self.lang.word2index['EOS']] + 
              [self.lang.word2index['PAD']]*(max_tgt_ln-len(tgts[i]))
              ) 
          tgts_padding_masks.append([False]*(len(tgts[i])+2) + [True]*(max_tgt_ln-len(tgts[i])))

        return (srcs_lns, np.array(converted_srcs), np.array(converted_tgts), np.array(srcs_padding_masks), np.array(tgts_padding_masks))

    def __len__(self):
        return self.batch_num


# train_dataset = HeadlinesDataset(pairs[1024:], lang, batch_size=config['batch_size'])
# test_dataset  = HeadlinesDataset(pairs[:1024], lang, batch_size=config['batch_size'])

# train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=4, pin_memory=True)
# test_loader = DataLoader(test_dataset, batch_size=1, shuffle=True, num_workers=4, pin_memory=True)

# src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(train_loader))
# src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(test_loader))
# print(src)
# print(tgt)
# print(src_key_padding_mask)
# print(tgt_key_padding_mask)

# + id="Lba3Xjzde5tp" colab_type="code" colab={}
class ABAEDataset(Dataset):
    def __init__(self, pairs, lang, m):
        self.m = m
        self.lang = lang
        self.pairs = pairs # [p[0] for p in pairs] + [p[1] for p in pairs]
        assert(self.m+1<len(self.pairs))

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

    def __getitem__(self, idx):
        ps = self.pairs[idx][0]
        ns_i = random.sample(range(len(self)), self.m)
        ns = [self.pairs[i][0] for i in ns_i]
        ps = self.lang.text2ids(ps)
        ns = [self.lang.text2ids(s) for s in ns]
        return ps, ns


def collate_fn(batch):
    ps_list = [b[0] for b in batch]
    ns_lists = [b[1] for b in batch]
    return ps_list, ns_lists


# ds = ABAEDataset(pairs, lang, 20)
# dataloader = DataLoader(ds, batch_size=3, collate_fn=collate_fn)
# btch = next(iter(dataloader))
# print(len(btch[0]))
# print(len(btch[1]))

# + [markdown] id="HTf3m8gdNcJr" colab_type="text"
# # Prepear word embeddings

# + id="umLFEfBM3Kkh" colab_type="code" colab={}
config = {
    'd_model': 512,
    'aspects_num': 100,
    'abae_regilarize': 0.000001,
    'batch_size': 64,
    'abae_negative_m': 20,
    'nhead': 4, # 8,
    'num_encoder_layers': 3, # 6
    'num_decoder_layers': 3, # 6
    'dim_feedforward': 2048,
    'pos_dropout': 0.2,
    'trans_dropout': 0.2,
    'n_warmup_steps': 4000,
}

# + id="Wty7UgXlNIXZ" colab_type="code" outputId="06b3214e-079e-4c36-b464-b95375abda18" colab={"base_uri": "https://localhost:8080/", "height": 128}
from random import shuffle

union = []
for pair in tqdm.tqdm(pairs):
  union.append(['SOS']+pair[0]+['EOS']+['PAD']+['UNK'])
  # union.append(['SOS']+pair[1]+['EOS']+['PAD']+['UNK'])

shuffle(union)


from gensim.models.word2vec import Word2Vec
from multiprocessing import cpu_count

model = Word2Vec(union, 
                 min_count = 0, 
                 workers=cpu_count(),
                 size=config['d_model'],
                 window=10,
                 negative=5,
                #  iter=100
                )


init_voc = set(lang.word2index.keys())
wv_voc = set(model.wv.vocab.keys())
print(len(wv_voc), len(init_voc))
print(wv_voc - init_voc)
print(init_voc - wv_voc)
# assert(len(init_voc)==len(wv_voc))


word_init_vectors = {}
for w in model.wv.vocab:
    word_init_vectors[w] = model[w].tolist()

# + [markdown] id="vYoF4ifEu-Z9" colab_type="text"
# # Prepear aspects vectors

# + id="FE2GPv2CGt_z" colab_type="code" outputId="4b7e743f-9bf3-4bd2-d381-759194d210c0" colab={"base_uri": "https://localhost:8080/", "height": 283}
lens = [lang.word2count[w] for w in lang.word2count]
print(len([w for w in lang.word2count if lang.word2count[w] > 1000 and lang.word2count[w] < 5000]))
sns.distplot([l for l in lens if l > 1000 and l < 10000])
plt.show()

# + id="On9n46RdvBSp" colab_type="code" outputId="6b918da0-5a6d-4bd7-d71f-ca75855208cc" colab={"base_uri": "https://localhost:8080/", "height": 1000}
X = []
words = []
for w in lang.word2count:
    if lang.word2count[w] > 1000 and lang.word2count[w] < 5000:
        X.append(word_init_vectors[w])
        words.append(w)
print(len(words))


# K means
from sklearn.cluster import KMeans
import numpy as np
kmeans = KMeans(n_clusters=config['aspects_num'], n_init=10, max_iter=500, random_state=42, verbose=0).fit(X)
av = kmeans.cluster_centers_.tolist()


def sim(word, aspect):
    return np.dot(word, aspect)

def sim2(word, aspect):
    return np.sum((word-aspect)**2)

neares_num = 100
for aspect_i in range(len(av)):
    aspect = np.array(av[aspect_i])
    similarities = []
    for word_i in range(len(X)):
        word = words[word_i]
        similarities.append([word, sim(X[word_i], aspect)])
    
    low_words = sorted(similarities, key=lambda x: x[1])[:neares_num]
    print(aspect_i, [w[0] for w in low_words])

# + [markdown] id="5h7Rvg6vWyYI" colab_type="text"
# # Model

# + id="Y4xDeTYojtoC" colab_type="code" colab={}
import math 
# Source: https://pytorch.org/tutorials/beginner/transformer_tutorial.html
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=200):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)


class ScheduledOptim():
    '''A simple wrapper class for learning rate scheduling'''

    def __init__(self, optimizer, d_model, n_warmup_steps):
        self._optimizer = optimizer
        self.n_warmup_steps = n_warmup_steps
        self.n_current_steps = 0
        self.init_lr = np.power(d_model, -0.5)

    def step(self):
        "Step with the inner optimizer"
        self._update_learning_rate()
        self._optimizer.step()

    def zero_grad(self):
        "Zero out the gradients by the inner optimizer"
        self._optimizer.zero_grad()

    def _get_lr_scale(self):
        return np.min([
            np.power(self.n_current_steps, -0.5),
            np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])

    def _update_learning_rate(self):
        ''' Learning rate scheduling per step '''

        self.n_current_steps += 1
        lr = self.init_lr * self._get_lr_scale()

        for param_group in self._optimizer.param_groups:
            param_group['lr'] = lr


def gen_nopeek_mask(length):
    mask = rearrange(torch.triu(torch.ones(length, length)) == 1, 'h w -> w h')
    mask = rearrange(torch.triu(torch.ones(length, length)) == 1, 'h w -> w h')
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask

class ConditionalHeadlineGenerator(nn.Module):
    def __init__(self, lang, device, vocab_size, max_seq_length, d_model, aspects_num, abae_regilarize, batch_size,  
                 abae_negative_m, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, pos_dropout, trans_dropout):
        super().__init__()
        # Constants
        self.lang = lang
        self.device = device
        self.d_model = d_model
        self.vocab_size = vocab_size
        self.aspects_num = aspects_num
        self.abae_regilarize = abae_regilarize
        self.abae_negative_m = abae_negative_m
        self.batch_size = batch_size

        self.max_seq_length = max_seq_length
        
        # auto .parameters()
        self.embed_tgt = nn.Embedding(vocab_size, d_model)
        self.embed_src = nn.Embedding(vocab_size, d_model)
        self.transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, trans_dropout)
        self.fc = nn.Linear(d_model, vocab_size)
        self.w_embs = self.embed_src # abae
        self.sentence_aspects = torch.nn.Linear(d_model, aspects_num) # abae

        # needs to be added to optimizer
        self.a_embs = torch.randn(aspects_num, d_model, requires_grad=True, device=device) # abae
        self.attention_M = torch.randn(d_model, d_model, requires_grad=True, device=device) # abae

        # other staff
        self.pos_enc = PositionalEncoding(d_model, pos_dropout, max_seq_length+1) # +1 for first conditional token
        self.criterion = nn.CrossEntropyLoss(ignore_index=self.lang.word2index['PAD'])
        self.optimizer = None

        #set device
        self.to(self.device)



    def tgt_emb(self, tgt):
        return self.embed_src(tgt)


    def src_emb(self, src): # for T and abae
        return self.embed_src(src)


    def make_sentence(self, embs):
        """
        input: embedded words
        output: weights
        """
        ys = torch.mean(embs, axis=0).unsqueeze(dim=1)
        # print(self.attention_M.shape) # torch.Size([300, 300])
        # print(ys.shape) # torch.Size([300, 1])
        # print(torch.mm(self.attention_M, ys).shape) # torch.Size([300, 1])
        # print(embs.shape) # torch.Size([19, 300])
        di = torch.mm(embs, torch.mm(self.attention_M, ys)) # [seq_len, dim] [[dim, dim] [dim, 1]]
        ai = torch.nn.functional.softmax(di, dim=1)
        zs = torch.mm(embs.T, ai)
        return zs


    def add_major_aspect(self, srcs_lns, src):
        """
        input: embedded sentences [Batch, seq len, dim]
        output: embedded words hstasked with aspect vectors [Batch, 1+seq len, dim]
        """
        av = torch.zeros(src.shape[0], 1, src.shape[2]).to(self.device) # [Batch, dim]
        for i in range(src.shape[0]):
          ys = torch.mean(src[i, :srcs_lns[i], :], axis=0).unsqueeze(dim=1) # [dim, 1]
          di = torch.mm(src[i, :srcs_lns[i], :], torch.mm(self.attention_M, ys)) # [seq_len, dim] [[dim, dim] [dim, 1]]
          ai = torch.nn.functional.softmax(di, dim=1) # [seq_len, 1]
          zs = torch.mm(src[i, :srcs_lns[i], :].T, ai) # [dim, 1]
          pt = torch.nn.functional.softmax(self.sentence_aspects(zs.T), dim=1).T # [apects num, 1]
          av[i, 0, :] = self.a_embs[torch.argmax(pt)]
        src = torch.cat([av, src], dim=1) # Batch, 1+seq len, dim]
        return src # [Batch, 1+seq len, dim]


    def init_embeddings(self, wv): # abae
        assert(self.embed_src.weight.data.shape == wv.shape)
        assert(self.embed_tgt.weight.data.shape == wv.shape)
        self.embed_src.weight.data = wv
        self.embed_tgt.weight.data = wv


    def make_sentence_raw(self, embs): # abae
        return torch.mean(embs, axis=0).unsqueeze(dim=1)


    def abae_forward(self, ps, ns_list):
        """
        ps - positive sequence. Tensor of shape (sentence_length,)
        ns_list - list of negaitve sequence. Tensor of shape [(sentence_length,), ...]
        """
        embs = self.w_embs(ps) # (sentence_length, emb_dim)
        zs = self.make_sentence(embs) # sentence from attentioned words vectors
        pt = torch.nn.functional.softmax(self.sentence_aspects(zs.T), dim=1).T
        rs = torch.mm(self.a_embs.T, pt) # sentence reconstraction by aspects vectors
        return rs, zs, [self.make_sentence_raw(self.w_embs(ns)) for ns in ns_list]


    def abae_step(self, b_ps, b_ns):
        """
        No grad zero here.
        Input:
        --------
        b_ps batch of positive sentences
        b_ns batch of lists of negative sentances
        """
        batch_loss = torch.Tensor([[0.0]]).to(self.device)
        for b_s in range(len(b_ps)):
            ps = torch.Tensor(b_ps[b_s]).long().to(self.device)
            ns_list = [torch.Tensor(ns).long().to(self.device) for ns in b_ns[b_s]]
            rs, ps, ns_list = self.abae_forward(ps, ns_list)
            sample_loss = torch.Tensor([[1.0]]).to(self.device)
            sample_loss -= torch.mm(rs.T, ps).to(self.device)
            for ns in ns_list:
                sample_loss += torch.mm(rs.T, ns).to(self.device)
            if sample_loss > 0:
                batch_loss += sample_loss
        batch_loss += self.abae_regilarize*torch.norm(torch.mm(self.a_embs, self.a_embs.T)-torch.eye(self.a_embs.shape[0]).to(self.device))
        return batch_loss


    def abae_train(self, pairs, epochs_n, progress_bar=False, logging=False, dataset_shuffle=False):
        """
        Be carefull ro turn your parameters into 
        train phase before launching this.
        """
        abae_dataset = ABAEDataset(pairs, self.lang, self.abae_negative_m)
        losses = []
        for epoch_i in range(epochs_n):
            abae_dataloader = DataLoader(abae_dataset, 
                                         batch_size=self.batch_size, 
                                         collate_fn=collate_fn,
                                         shuffle=dataset_shuffle)
            epoch_loss = torch.Tensor([[0.0]]).to(self.device)
            for b_ps, b_ns in tqdm.tqdm(abae_dataloader, disable=(not progress_bar)):
                self.optimizer.zero_grad()
                batch_loss = self.abae_step(b_ps, b_ns)
                epoch_loss += batch_loss.item()
                batch_loss.backward()
                self.optimizer.step()          
            losses.append((epoch_i, epoch_loss.item()))
            if logging:
                print(f"ABAE train epoch: {epoch_i} - loss: {epoch_loss}")
        
        return losses

    def transformer_train(self, pairs, epochs_n, progress_bar=False, logging=False, dataset_shuffle=False):
        """
        Transformer train
        """
        train_size = int(0.8*len(pairs))
        train_T_dataset = HeadlinesDataset(pairs[:train_size], self.lang, batch_size=self.batch_size)
        test_T_dataset  = HeadlinesDataset(pairs[train_size:], self.lang, batch_size=self.batch_size)
        
        train_losses = []
        test_losses = []
        for epoch_i in range(epochs_n):
            train_T_loader = DataLoader(train_T_dataset, batch_size=1, shuffle=dataset_shuffle, num_workers=0, pin_memory=True)
            test_T_loader = DataLoader(test_T_dataset, batch_size=1, shuffle=dataset_shuffle, num_workers=0, pin_memory=True)

            self.train()
            train_epoch_loss = 0.0
            for srcs_lns, src, tgt, src_key_padding_mask, tgt_key_padding_mask in tqdm.tqdm(train_T_loader, disable=(not progress_bar)):
                self.optimizer.zero_grad()
                batch_loss = torch.Tensor([[0.0]]).to(self.device)

                src, src_key_padding_mask = src[0].to(device), src_key_padding_mask[0].to(device)
                tgt, tgt_key_padding_mask = tgt[0].to(device), tgt_key_padding_mask[0].to(device)
                memory_key_padding_mask = src_key_padding_mask.clone()
                tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]
                tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]
                tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device

                output = self.transformer_forward(srcs_lns, src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)
                batch_loss += self.criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))
                batch_loss.backward()
                self.optimizer.step()
                train_epoch_loss += batch_loss.item()
            train_losses.append((epoch_i, train_epoch_loss))


            self.eval()
            test_epoch_loss = 0.0
            for srcs_lns, src, tgt, src_key_padding_mask, tgt_key_padding_mask in tqdm.tqdm(test_T_loader, disable=(not progress_bar)):
                with torch.no_grad():
                    src, src_key_padding_mask = src[0].to(device), src_key_padding_mask[0].to(device) # to device
                    tgt, tgt_key_padding_mask = tgt[0].to(device), tgt_key_padding_mask[0].to(device) # to device
                    memory_key_padding_mask = src_key_padding_mask.clone()
                    tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]
                    tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]
                    tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device

                    output = self.transformer_forward(srcs_lns, src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)
                    loss = self.criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))
                    test_epoch_loss += loss.item()
            test_losses.append((epoch_i, test_epoch_loss))

            print(f"\nEpoch: {epoch_i} Train loss: {train_epoch_loss} Test loss: {test_epoch_loss}")
        
        return train_losses, test_losses

    def full_train(self, pairs, epochs_n, progress_bar=False, logging=False, dataset_shuffle=False):
        """
        Transformer and abae simultaneous train
        """
        train_T_dataset = HeadlinesDataset(pairs, self.lang, batch_size=self.batch_size)
        test_T_dataset  = HeadlinesDataset(pairs, self.lang, batch_size=self.batch_size)
        train_abae_dataset = ABAEDataset(pairs, self.lang, self.abae_negative_m)
        
        train_losses = []
        test_losses = []
        for epoch_i in range(epochs_n):
            train_T_loader = DataLoader(train_T_dataset, batch_size=1, shuffle=dataset_shuffle, num_workers=0, pin_memory=True)
            test_T_loader = DataLoader(test_T_dataset, batch_size=1, shuffle=dataset_shuffle, num_workers=0, pin_memory=True)
            train_abae_loader = DataLoader(train_abae_dataset, batch_size=self.batch_size, collate_fn=collate_fn)
            abae_iterator = iter(train_abae_loader)

            self.train()
            train_epoch_loss = 0.0
            for srcs_lns, src, tgt, src_key_padding_mask, tgt_key_padding_mask in tqdm.tqdm(train_T_loader, disable=(not progress_bar)):
                self.optimizer.zero_grad()
                batch_loss = torch.Tensor([[0.0]]).to(self.device)

                b_ps, b_ns = next(abae_iterator)
                batch_loss += self.abae_step(b_ps, b_ns)

                src, src_key_padding_mask = src[0].to(device), src_key_padding_mask[0].to(device)
                tgt, tgt_key_padding_mask = tgt[0].to(device), tgt_key_padding_mask[0].to(device)
                memory_key_padding_mask = src_key_padding_mask.clone()
                tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]
                tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]
                tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device

                output = self.transformer_forward(srcs_lns, src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)
                batch_loss += self.criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))
                batch_loss.backward()
                self.optimizer.step()
                train_epoch_loss += batch_loss.item()
            train_losses.append((epoch_i, train_epoch_loss))


            self.eval()
            test_epoch_loss = 0.0
            for srcs_lns, src, tgt, src_key_padding_mask, tgt_key_padding_mask in tqdm.tqdm(test_T_loader, disable=(not progress_bar)):
                with torch.no_grad():
                    src, src_key_padding_mask = src[0].to(device), src_key_padding_mask[0].to(device) # to device
                    tgt, tgt_key_padding_mask = tgt[0].to(device), tgt_key_padding_mask[0].to(device) # to device
                    memory_key_padding_mask = src_key_padding_mask.clone()
                    tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]
                    tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]
                    tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device

                    output = self.transformer_forward(srcs_lns, src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)
                    loss = self.criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))
                    test_epoch_loss += loss.item()
            test_losses.append((epoch_i, test_epoch_loss))

            print(f"\nEpoch: {epoch_i} Train loss: {train_epoch_loss} Test loss: {test_epoch_loss}")
        
        return train_losses, test_losses


    def transformer_forward(self, srcs_lns, src, tgt, src_key_padding_mask, tgt_key_padding_mask, memory_key_padding_mask, tgt_mask):
        tgt = tgt.permute(1, 0)
        src = self.embed_src(src)
        src = self.add_major_aspect(srcs_lns, src)
        src = src.permute(1, 0, 2)
        src = self.pos_enc(src*math.sqrt(self.d_model)) # why every time multiply? just init it somewhere
        tgt = self.pos_enc(self.embed_tgt(tgt) * math.sqrt(self.d_model))
        output = self.transformer(src, tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask,
                                  tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask)
        output = output.permute(1, 0, 2)
        return self.fc(output)


    def set_optimizer(self, Optimizer, **optimizer_args):
        self.optimizer = Optimizer(self.get_parameters(), **optimizer_args)

    def get_parameters(self):
        return list(self.parameters())+[self.a_embs, self.attention_M]

    def get_aspects_vector(self, inp):
        pass


    def transformer_train_step(self, inp):
        aspects_vector = self.get_aspects_vector(inp)
        pass

    def init_embs(self, wv, av):
        pass





model = ConditionalHeadlineGenerator(lang, device, lang.n_words, lang.max_seq_length,
                            config['d_model'], config['aspects_num'],
                            config['abae_regilarize'], config['batch_size'],
                            config['abae_negative_m'],
                            config['nhead'], config['num_encoder_layers'],
                            config['num_decoder_layers'], config['dim_feedforward'],  
                            config['pos_dropout'], config['trans_dropout'])


wv = torch.randn(model.vocab_size, model.d_model)
for w in word_init_vectors.keys():
    ind = lang.word2index[w]
    wv[ind, :] = torch.Tensor(word_init_vectors[w])
av = torch.Tensor(av).shape
model.init_embs(wv, av)

# + [markdown] id="dDiRWhUarIgq" colab_type="text"
# # Abae pretrain

# + id="r0xrogLEqCKm" colab_type="code" colab={}
model.set_optimizer(torch.optim.Adam, lr=0.001)

# + id="sVB0U9h09sG2" colab_type="code" outputId="a7f78dbb-74bc-401d-81ae-0d694c0056e4" colab={"base_uri": "https://localhost:8080/", "height": 126}
for param in model.embed_src.parameters():
    param.requires_grad = False
# print(model.embed_src.weight.requires_grad)
losses = model.abae_train(pairs, 3, progress_bar=True, logging=True)

# + id="qBDn1mVdv6_2" colab_type="code" outputId="9f84923f-d081-4663-86e2-cbcf9161f10e" colab={"base_uri": "https://localhost:8080/", "height": 72}
for param in model.embed_src.parameters():
    param.requires_grad = True
# print(model.embed_src.weight.requires_grad)
losses = model.abae_train(pairs, 3, progress_bar=True, logging=True)

# + [markdown] id="dimisX_jrLJy" colab_type="text"
# # Transformer train

# + id="UP-wlVQ9qF28" colab_type="code" colab={}
model.optimizer = ScheduledOptim(
                      optimizer=torch.optim.Adam(model.get_parameters(), betas=(0.9, 0.98), eps=1e-09),
                      d_model=config['d_model'], 
                      n_warmup_steps=config['n_warmup_steps']
                    )

# + id="h_aAcemy9sQD" colab_type="code" outputId="1e202f98-426d-441c-9e34-99c80248e190" colab={"base_uri": "https://localhost:8080/", "height": 235}
for param in model.embed_src.parameters():
    param.requires_grad = False
# model.embed_src.eval()
# print(model.embed_src.weight.requires_grad)
losses = model.full_train(pairs, 3, progress_bar=True, logging=True)

# + id="m-qdpltW9shD" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 235} outputId="5420091a-c593-41c7-9fac-7c17d54f1974"
for param in model.embed_src.parameters():
    param.requires_grad = True
# model.embed_src.eval()
# print(model.embed_src.weight.requires_grad)
losses = model.full_train(pairs, 3, progress_bar=True, logging=True)

# + id="8i7YvzAgWlbc" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 744} outputId="25decfdd-1c17-4c47-bd20-6c693e9ef810"
# for param in model.embed_src.parameters():
#     param.requires_grad = True
# model.embed_src.eval()
# print(model.embed_src.weight.requires_grad)
losses = model.transformer_train(pairs, 10, progress_bar=True, logging=True)

# + [markdown] id="vppAOhOzxzN3" colab_type="text"
# # Aspect inference

# + id="eRfY_rbsULY6" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 1000} outputId="5e7702ce-b688-4db8-a071-341fcde4126a"
aspects_vectors = model.a_embs.data.tolist()
words_vectors   = model.embed_src.weight.data.tolist()

words = []
word_v = []
for w in model.lang.word2count:
    if lang.word2count[w] > 1000 and lang.word2count[w] < 5000:
        word_v.append(words_vectors[model.lang.word2index[w]])
        words.append(w)
print(len(words))



def sim(word, aspect):
    return np.dot(word, aspect)

def sim2(word, aspect):
    return np.sum((word-aspect)**2)

neares_num = 100
for aspect_i in range(len(aspects_vectors)):
    aspect = np.array(aspects_vectors[aspect_i])
    similarities = []
    for word_i in range(len(word_v)):
        word = words[word_i]
        similarities.append([word, sim(word_v[word_i], aspect)])
    
    low_words = sorted(similarities, key=lambda x: x[1])[:neares_num]
    print(aspect_i, [w[0] for w in low_words])


# + [markdown] id="pjiZdqPkx3jZ" colab_type="text"
# # Generator inference

# + id="VJivaP2o8357" colab_type="code" colab={}
def forced_inference(model, src_ln, src, tgt, src_key_padding_mask, tgt_key_padding_mask):
  # TODO: add accuracy / recall 
  memory_key_padding_mask = src_key_padding_mask.clone()  
  tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]
  tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]
  tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device
  # print(src.shape)
  # print(tgt_inp.shape)
  # print(tgt_out.shape)
  # print(tgt_key_padding_mask_inp.shape)
  # print(tgt_mask.shape)
  # print()
  output = model.transformer_forward(src_ln, src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)
  tokens = torch.argmax(output[0], axis=1)
  return tokens

def greedy_inference(model, src_ln, src, tgt, src_key_padding_mask, tgt_key_padding_mask, lang):
  memory_key_padding_mask = src_key_padding_mask.clone() 
  tgt_inp, tgt_out = tgt[:, :-1].clone().detach(), tgt[:, 1:].clone().detach()
  tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]
  tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device

  tokens = [tgt_inp[0,0].item()]
  tgt_inp_ = tgt_inp[:]
  
  for tgt_i in range(tgt.shape[1]-1):
    tt = torch.Tensor(tokens)
    tgt_inp_[:, :tt.shape[0]] = tt[:]
    output = model.transformer_forward(src_ln, src, tgt_inp_, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)
    next_token = torch.argmax(output[0][tt.shape[0]+1]).item()
    tokens.append(next_token)
    if next_token == lang.word2index['EOS']:
      break
  return torch.Tensor(tokens[1:])

def inference(model, src_ln, src, tgt, src_key_padding_mask, tgt_key_padding_mask, lang, mode):
  """
  src: Tensor [1, S]
  tgt: Tensor [1, T]
  modes: forced / greedy [/ beam search]
  """
  model.eval()

  if mode == "forced":
    tokens = forced_inference(model, src_ln, src, tgt, src_key_padding_mask, tgt_key_padding_mask)
    
  elif mode == "greedy":
    tokens = greedy_inference(model, src_ln, src, tgt, src_key_padding_mask, tgt_key_padding_mask, lang)

  print(f"Mode: {mode}")
  print(f'Source: {lang.tensor2text(src[0])}')
  print(f'Target: {lang.tensor2text(tgt[0])}')
  print(f'Generated: {lang.tensor2text(tokens)}')


# + id="cbCXQOP0oyh7" colab_type="code" colab={}
train_T_dataset = HeadlinesDataset(pairs[:9000], model.lang, batch_size=model.batch_size)
test_T_dataset  = HeadlinesDataset(pairs[9000:10000], model.lang, batch_size=model.batch_size)
train_T_loader = DataLoader(train_T_dataset, batch_size=1, shuffle=True, num_workers=0, pin_memory=True)
test_T_loader = DataLoader(test_T_dataset, batch_size=1, shuffle=True, num_workers=0, pin_memory=True)

# + id="tAwc2C0qVeAK" colab_type="code" outputId="447138b9-1519-4eb2-e4ba-17cef064c3c0" colab={"base_uri": "https://localhost:8080/", "height": 182}
srcs_lns, src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(train_T_loader))
text_i = 0
src_ln = srcs_lns[text_i:text_i+1]
src, src_key_padding_mask = src[0][text_i:text_i+1].to(device), src_key_padding_mask[0][text_i:text_i+1].to(device) # to device
tgt, tgt_key_padding_mask = tgt[0][text_i:text_i+1].to(device), tgt_key_padding_mask[0][text_i:text_i+1].to(device) # to device
inference(model, src_ln, src, tgt, src_key_padding_mask, tgt_key_padding_mask, lang, 'forced')
inference(model, src_ln, src, tgt, src_key_padding_mask, tgt_key_padding_mask, lang, 'greedy')

# + id="3JaS9PPVKg6I" colab_type="code" outputId="fc161764-22e1-4168-feaa-8e24cbfb4f73" colab={"base_uri": "https://localhost:8080/", "height": 182}
srcs_lns, src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(test_T_loader))
text_i = 0
src_ln = srcs_lns[text_i:text_i+1]
src, src_key_padding_mask = src[0][text_i:text_i+1].to(device), src_key_padding_mask[0][text_i:text_i+1].to(device) # to device
tgt, tgt_key_padding_mask = tgt[0][text_i:text_i+1].to(device), tgt_key_padding_mask[0][text_i:text_i+1].to(device) # to device
inference(model, src_ln, src, tgt, src_key_padding_mask, tgt_key_padding_mask, lang, 'forced')
inference(model, src_ln, src, tgt, src_key_padding_mask, tgt_key_padding_mask, lang, 'greedy')

# + id="uVKK1VPi3BI5" colab_type="code" colab={}


# + id="LWjuMGMe3Bg_" colab_type="code" colab={}


# + id="bOwgpq0O3BeL" colab_type="code" colab={}


# + id="ItotVW3E3BbN" colab_type="code" colab={}


# + id="KHldzZGu3BUD" colab_type="code" colab={}


# + id="nRZLACVe3BRI" colab_type="code" colab={}


# + id="Dmjdij1wKhCV" colab_type="code" colab={}


# + id="QMZ-3RmlVc6x" colab_type="code" colab={}


# + id="oeMGFrfczXBd" colab_type="code" outputId="d2e8ab62-4ba6-4612-a284-2d662163356c" colab={"base_uri": "https://localhost:8080/", "height": 181}
# failed sanity check

src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(train_loader))
text_i = 3
src, src_key_padding_mask = src[0][text_i:text_i+1].to(device), src_key_padding_mask[0][text_i:text_i+1].to(device) # to device
tgt, tgt_key_padding_mask = tgt[0][text_i:text_i+1].to(device), tgt_key_padding_mask[0][text_i:text_i+1].to(device) # to device

tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]
tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]
tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device
print(src.shape)
print(tgt_inp.shape)
print(tgt_out.shape)
print(tgt_key_padding_mask_inp.shape)
print(tgt_mask.shape)
print()

# tgt_key_padding_mask_inp = torch.Tensor([[False]+[True]*(tgt_inp.shape[1]-1)]).bool().to(device)
model.eval()
tgt_mask[:, 1:] = tgt_mask[0, 1] # sanity check
output = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)
print(output.shape)
print(output[0, :, 9]) # they should be same

# + id="TAaTFU_ENoVu" colab_type="code" outputId="545811f1-979d-4fab-c8d1-f88758b514eb" colab={"base_uri": "https://localhost:8080/", "height": 53}


# + id="ET5qGZpDN7Wt" colab_type="code" outputId="3d9a1580-73b6-4101-e2cb-5c45e1cad7da" colab={"base_uri": "https://localhost:8080/", "height": 162}
output = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)
print(output)

# + id="y8FiTkKpzLiQ" colab_type="code" outputId="985523e9-5d5c-4065-9911-80c73d8fb199" colab={"base_uri": "https://localhost:8080/", "height": 108}
print(tgt_inp)
print(tgt_key_padding_mask_inp)
print()

# + id="fld_JsGDzR84" colab_type="code" colab={}

loss = criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))

# + id="EqhVe5f084Gj" colab_type="code" colab={}
# inference


src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(test_loader))
print(src)
print(tgt)
print(src_key_padding_mask)
print(tgt_key_padding_mask)

# + id="WAOvRJG184Cp" colab_type="code" colab={}


# + id="eI-3Spl-0gls" colab_type="code" colab={}


# + id="nRDpuv3ezUT_" colab_type="code" colab={}


# + id="5MYZdOtCzUoi" colab_type="code" colab={}


# + id="MCJHC_udqSww" colab_type="code" outputId="0590fa2f-b2c3-485a-c32d-ae674d898949" colab={"base_uri": "https://localhost:8080/", "height": 243}
src_, src_key_padding_mask_ = src[0].to(device), src_key_padding_mask[0].to(device)
tgt_, tgt_key_padding_mask_ = tgt[0].to(device), tgt_key_padding_mask[0].to(device)
memory_key_padding_mask = src_key_padding_mask_.clone()
tgt_inp, tgt_out = tgt_[:, :-1], tgt_[:, 1:]
tgt_key_padding_mask_inp = tgt_key_padding_mask_[:, :-1]
tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device)

# + id="wvaPZRRvuOg6" colab_type="code" colab={}
print(tgt_key_padding_mask_.shape)
print(src_.shape)
print(tgt_inp.shape)
print(src_key_padding_mask_.shape)
print(memory_key_padding_mask.shape)
print(tgt_mask.shape)

# + id="Mlsp_16Etdii" colab_type="code" colab={}
output = model(src_, tgt_inp, src_key_padding_mask_, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)

# + id="mTX_erk1t2WD" colab_type="code" colab={}
 torch.isnan(output).any()

# + id="b5_9Agftv8Vz" colab_type="code" colab={}
src_.shape

# + id="qbyDvyKzwQyR" colab_type="code" colab={}
src_[0][0]

# + id="xwtgvBOdw6Ft" colab_type="code" colab={}
src[0][9]

# + id="ddF9kC1Uw0B9" colab_type="code" colab={}
lang.tensor2text(tgt[0][9])

# + id="Wd-JydO0vuCn" colab_type="code" colab={}
for row_i in range(src_.shape[0]):
  print(row_i, src_[row_i][1])

# + id="CEKnVJ-fp6Eh" colab_type="code" colab={}
src, src_key_padding_mask = src[0].to(device), src_key_padding_mask[0].to(device) # to device
tgt, tgt_key_padding_mask = tgt[0].to(device), tgt_key_padding_mask[0].to(device) # to device
memory_key_padding_mask = src_key_padding_mask.clone()
tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]
tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]
tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device

optimizer.zero_grad()
output = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)
loss = criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))
loss.backward()
optimizer.step()
train_loss += loss.item()

# + id="MiETl8cebCGq" colab_type="code" colab={}
output = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask_inp, memory_key_padding_mask, tgt_mask)
criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))

# + id="QrOG2X88b8XL" colab_type="code" colab={}
src_ = src.permute(1, 0)
tgt_ = tgt_inp.permute(1, 0)
# src_ = model.pos_enc(model.embed_src(src_) * math.sqrt(model.d_model)) # why every time multiply? just init it somewhere
# tgt_ = model.pos_enc(model.embed_tgt(tgt_) * math.sqrt(model.d_model))

# + id="vMlRQSWTcT7p" colab_type="code" colab={}
src_[:, 0].unsqueeze(0)

# + id="KZd99ArrcC57" colab_type="code" colab={}
print(model.embed_src(src_[:, 0].unsqueeze(0)))

# + id="8ei_HEKSelPz" colab_type="code" colab={}
model.embed_src(torch.Tensor([[lang.word2index['PAD']]]).long().to(device) )

# + id="8UklVda-bndB" colab_type="code" colab={}


output = self.transformer(src, tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask,
                          tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask)
output = output.permute(1, 0, 2)
return self.fc(output)

# + id="03qw5a46bieg" colab_type="code" colab={}
output

# + id="bcwjAlMA1y7Y" colab_type="code" colab={}

global_embeddings = nn.Embedding(lang.n_words, config['d_model'])
wv = torch.Tensor([word_init_vectors[lang.index2word[idx]] for idx in lang.index2word.keys()])
global_embeddings.weight.data = wv


criterion = nn.CrossEntropyLoss(ignore_index=lang.word2index['PAD'])

d_model = config['d_model']
pos_dropout = config['pos_dropout']
nhead = config['nhead']
num_encoder_layers = config['num_encoder_layers']
num_decoder_layers = config['num_decoder_layers']
dim_feedforward = config['dim_feedforward']
trans_dropout = config['trans_dropout']



embed_src = global_embeddings
embed_tgt = global_embeddings
pos_enc = PositionalEncoding(d_model, pos_dropout, lang.max_seq_length)
transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, trans_dropout)
fc = nn.Linear(d_model, lang.n_words)

# + id="ZHYTLN6Q55Bd" colab_type="code" colab={}
src, tgt, src_key_padding_mask, tgt_key_padding_mask =  next(iter(train_loader))
src, src_key_padding_mask = src[0], src_key_padding_mask[0].to(device) # to device
tgt, tgt_key_padding_mask = tgt[0], tgt_key_padding_mask[0].to(device) # to device
memory_key_padding_mask = src_key_padding_mask.clone()
tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]
tgt_key_padding_mask_inp = tgt_key_padding_mask[:, :-1]
tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to(device) # also to device

print(src.shape)
print(tgt.shape)
print(src_key_padding_mask.shape)
print(tgt_key_padding_mask.shape)
print(tgt_inp.shape)
print(tgt_out.shape)
print(tgt_mask.shape)

# + id="BGmF1vKh7dCT" colab_type="code" colab={}
src = src.permute(1, 0)
tgt_inp = tgt_inp.permute(1, 0)
src = pos_enc(embed_src(src) * math.sqrt(d_model))
tgt_inp = pos_enc(embed_tgt(tgt_inp) * math.sqrt(d_model))

# + id="f1ovGU5J90GA" colab_type="code" colab={}
print(src.shape, tgt_inp.shape)

# + id="xtXVOdxY-lfP" colab_type="code" colab={}
print(src.shape)
print(tgt_inp.shape)
print(tgt_mask.shape)
print(src_key_padding_mask.shape)
print(tgt_key_padding_mask_inp.shape)
print(memory_key_padding_mask.shape)

# + id="It8h9lga3Pii" colab_type="code" colab={}
output = transformer(src, tgt_inp, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask,
                          tgt_key_padding_mask=tgt_key_padding_mask_inp, memory_key_padding_mask=memory_key_padding_mask)
output = output.permute(1, 0, 2)
output = fc(output)

# + id="L0-s6h7X9te1" colab_type="code" colab={}
loss = criterion(rearrange(output, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))

# + id="a_5b9bxKAJSN" colab_type="code" colab={}


# + id="C0LtWb5iW1A8" colab_type="code" colab={}
import math

import torch
from torch import nn


class LanguageTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, max_seq_length, pos_dropout, trans_dropout):
        super().__init__()
        self.d_model = d_model
        self.embed_src = nn.Embedding(vocab_size, d_model)
        self.embed_tgt = nn.Embedding(vocab_size, d_model)
        self.pos_enc = PositionalEncoding(d_model, pos_dropout, max_seq_length)

        self.transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, trans_dropout)
        self.fc = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt, src_key_padding_mask, tgt_key_padding_mask, memory_key_padding_mask, tgt_mask):
        # src = rearrange(src, 'n s -> s n')
        src = src.permute(1, 0)
        # tgt = rearrange(tgt, 'n t -> t n')
        tgt = tgt.permute(1, 0)
        src = self.pos_enc(self.embed_src(src) * math.sqrt(self.d_model)) # why every time multiply? just init it somewhere
        tgt = self.pos_enc(self.embed_tgt(tgt) * math.sqrt(self.d_model))

        output = self.transformer(src, tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask,
                                  tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask)
        # output = rearrange(output, 't n e -> n t e')
        output = output.permute(1, 0, 2)
        return self.fc(output)


# + id="Unw4BYAFZAVX" colab_type="code" colab={}
        total_step += 1

        src, src_key_padding_mask = src[0].to('cuda'), src_key_padding_mask[0].to('cuda')
        tgt, tgt_key_padding_mask = tgt[0].to('cuda'), tgt_key_padding_mask[0].to('cuda')

        memory_key_padding_mask = src_key_padding_mask.clone()
        tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]
        tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to('cuda')

        optim.zero_grad()
        outputs = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask[:, :-1], memory_key_padding_mask, tgt_mask)
        loss = criterion(rearrange(outputs, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))

        loss.backward()
        optim.step_and_update_lr()

        total_loss += loss.item()
        train_losses.append((step, loss.item()))
        pbar.update(1)
        if step % print_every == print_every - 1:
            pbar.close()
            print(f'Epoch [{epoch + 1} / {num_epochs}] \t Step [{step + 1} / {len(train_loader)}] \t '
                  f'Train Loss: {total_loss / print_every}')
            total_loss = 0

            pbar = tqdm(total=print_every, leave=False)

    pbar.close()
    val_loss = validate(valid_loader, model, criterion)
    val_losses.append((total_step, val_loss))
    if val_loss < lowest_val:
        lowest_val = val_loss
        torch.save(model, 'output/transformer.pth')
    print(f'Val Loss: {val_loss}')

# + id="l5m8S0o8ZAiG" colab_type="code" colab={}


# + id="f91eDhSkXFFD" colab_type="code" colab={}


# + id="CWMTJGMVeO0w" colab_type="code" colab={}
import pickle
import random
import numpy as np
from torch.utils.data import Dataset


class ParallelLanguageDataset(Dataset):
    def __init__(self, data_path_1, data_path_2, num_tokens, max_seq_length):
        self.num_tokens = num_tokens
        self.data_1, self.data_2, self.data_lengths = load_data(data_path_1, data_path_2, max_seq_length)

        self.batches = gen_batches(num_tokens, self.data_lengths)

    def __getitem__(self, idx):
        src, src_mask = getitem(idx, self.data_1, self.batches, True)
        tgt, tgt_mask = getitem(idx, self.data_2, self.batches, False)

        return src, src_mask, tgt, tgt_mask

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

    def shuffle_batches(self):
        self.batches = gen_batches(self.num_tokens, self.data_lengths)


def gen_batches(num_tokens, data_lengths):
    # Shuffle all the indices
    for k, v in data_lengths.items():
        random.shuffle(v)

    batches = []
    prev_tokens_in_batch = 1e10
    for k in sorted(data_lengths):
        v = data_lengths[k]
        total_tokens = (k[0] + k[1]) * len(v)

        while total_tokens > 0:
            tokens_in_batch = min(total_tokens, num_tokens) - min(total_tokens, num_tokens) % (k[0] + k[1])
            sentences_in_batch = tokens_in_batch // (k[0] + k[1])

            # Combine with previous batch?
            if tokens_in_batch + prev_tokens_in_batch <= num_tokens:
                batches[-1].extend(v[:sentences_in_batch])
                prev_tokens_in_batch += tokens_in_batch
            else:
                batches.append(v[:sentences_in_batch])
                prev_tokens_in_batch = tokens_in_batch
            v = v[sentences_in_batch:]

            total_tokens = (k[0] + k[1]) * len(v)
    return batches


def load_data(data_path_1, data_path_2, max_seq_length):
    with open(data_path_1, 'rb') as f:
        data_1 = pickle.load(f)
    with open(data_path_2, 'rb') as f:
        data_2 = pickle.load(f)

    data_lengths = {}
    for i, (str_1, str_2) in enumerate(zip(data_1, data_2)):
        if 0 < len(str_1) <= max_seq_length and 0 < len(str_2) <= max_seq_length - 2:
            if (len(str_1), len(str_2)) in data_lengths:
                data_lengths[(len(str_1), len(str_2))].append(i)
            else:
                data_lengths[(len(str_1), len(str_2))] = [i]
    return data_1, data_2, data_lengths


def getitem(idx, data, batches, src):
    sentence_indices = batches[idx]
    if src:
        batch = [data[i] for i in sentence_indices]
    else:
        batch = [[2] + data[i] + [3] for i in sentence_indices]

    seq_length = 0
    for sentence in batch:
        if len(sentence) > seq_length:
            seq_length = len(sentence)

    masks = []
    for i, sentence in enumerate(batch):
        masks.append([False for _ in range(len(sentence))] + [True for _ in range(seq_length - len(sentence))])
        batch[i] = sentence + [0 for _ in range(seq_length - len(sentence))]

    return np.array(batch), np.array(masks)


# + id="e4IQmfdDeCPi" colab_type="code" colab={}
class ScheduledOptim():
    '''A simple wrapper class for learning rate scheduling'''

    def __init__(self, optimizer, d_model, n_warmup_steps):
        self._optimizer = optimizer
        self.n_warmup_steps = n_warmup_steps
        self.n_current_steps = 0
        self.init_lr = np.power(d_model, -0.5)

    def step_and_update_lr(self):
        "Step with the inner optimizer"
        self._update_learning_rate()
        self._optimizer.step()

    def zero_grad(self):
        "Zero out the gradients by the inner optimizer"
        self._optimizer.zero_grad()

    def _get_lr_scale(self):
        return np.min([
            np.power(self.n_current_steps, -0.5),
            np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])

    def _update_learning_rate(self):
        ''' Learning rate scheduling per step '''

        self.n_current_steps += 1
        lr = self.init_lr * self._get_lr_scale()

        for param_group in self._optimizer.param_groups:
            param_group['lr'] = lr



# + id="NV2jC2Isd8H1" colab_type="code" colab={}
def train(train_loader, valid_loader, model, optim, criterion, num_epochs):
    print_every = 500
    model.train()

    lowest_val = 1e9
    train_losses = []
    val_losses = []
    total_step = 0
    for epoch in range(num_epochs):
        pbar = tqdm(total=print_every, leave=False)
        total_loss = 0

        train_loader.dataset.shuffle_batches()
        for step, (src, src_key_padding_mask, tgt, tgt_key_padding_mask) in enumerate(iter(train_loader)):
            total_step += 1

            src, src_key_padding_mask = src[0].to('cuda'), src_key_padding_mask[0].to('cuda')
            tgt, tgt_key_padding_mask = tgt[0].to('cuda'), tgt_key_padding_mask[0].to('cuda')

            memory_key_padding_mask = src_key_padding_mask.clone()
            tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:]
            tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to('cuda')

            optim.zero_grad()
            outputs = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask[:, :-1], memory_key_padding_mask, tgt_mask)
            loss = criterion(rearrange(outputs, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))

            loss.backward()
            optim.step_and_update_lr()

            total_loss += loss.item()
            train_losses.append((step, loss.item()))
            pbar.update(1)
            if step % print_every == print_every - 1:
                pbar.close()
                print(f'Epoch [{epoch + 1} / {num_epochs}] \t Step [{step + 1} / {len(train_loader)}] \t '
                      f'Train Loss: {total_loss / print_every}')
                total_loss = 0

                pbar = tqdm(total=print_every, leave=False)

        pbar.close()
        val_loss = validate(valid_loader, model, criterion)
        val_losses.append((total_step, val_loss))
        if val_loss < lowest_val:
            lowest_val = val_loss
            torch.save(model, 'output/transformer.pth')
        print(f'Val Loss: {val_loss}')
    return train_losses, val_losses


def validate(valid_loader, model, criterion):
    pbar = tqdm(total=len(iter(valid_loader)), leave=False)
    model.eval()

    total_loss = 0
    for src, src_key_padding_mask, tgt, tgt_key_padding_mask in iter(valid_loader):
        with torch.no_grad():
            src, src_key_padding_mask = src[0].to('cuda'), src_key_padding_mask[0].to('cuda')
            tgt, tgt_key_padding_mask = tgt[0].to('cuda'), tgt_key_padding_mask[0].to('cuda')
            memory_key_padding_mask = src_key_padding_mask.clone()
            tgt_inp = tgt[:, :-1]
            tgt_out = tgt[:, 1:].contiguous()
            tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to('cuda')

            outputs = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask[:, :-1], memory_key_padding_mask, tgt_mask)
            loss = criterion(rearrange(outputs, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))

            total_loss += loss.item()
            pbar.update(1)

    pbar.close()
    model.train()
    return total_loss / len(valid_loader)


def gen_nopeek_mask(length):
    mask = rearrange(torch.triu(torch.ones(length, length)) == 1, 'h w -> w h')
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask


# + id="MwUsgjTygwm-" colab_type="code" colab={}
kwargs = {
    'num_epochs': 20,
    'max_seq_length': 96,
    'num_tokens': 2000,
    'vocab_size': 10000+4,
    'd_model': 512,
    'num_encoder_layers': 6,
    'num_decoder_layers': 6,
    'dim_feedforward': 2048,
    'nhead': 8,
    'pos_dropout': 0.1,
    'trans_dropout': 0.1,
    'n_warmup_steps': 4000
}

# + id="RspHKWOthXrs" colab_type="code" colab={}
train_dataset = ParallelLanguageDataset(root + 'data/processed/en/train.pkl',
                                        root + 'data/processed/fr/train.pkl',
                                        kwargs['num_tokens'], kwargs['max_seq_length'])
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=4, pin_memory=True)
valid_dataset = ParallelLanguageDataset(root + 'data/processed/en/val.pkl',
                                        root + 'data/processed/fr/val.pkl',
                                        kwargs['num_tokens'], kwargs['max_seq_length'])
valid_loader = DataLoader(valid_dataset, batch_size=1, shuffle=True, num_workers=4, pin_memory=True)

# + id="jnP2tyOGhgMO" colab_type="code" colab={}
src, src_key_padding_mask, tgt, tgt_key_padding_mask = train_dataset[100]
print(src.shape)
print(tgt.shape)

# + id="3S7ltY2jrR2m" colab_type="code" colab={}


# + id="xnA9YQ9Mgvei" colab_type="code" colab={}
model = LanguageTransformer(kwargs['vocab_size'], kwargs['d_model'], kwargs['nhead'], kwargs['num_encoder_layers'],
                            kwargs['num_decoder_layers'], kwargs['dim_feedforward'], kwargs['max_seq_length'],
                            kwargs['pos_dropout'], kwargs['trans_dropout']).to('cuda')
for p in model.parameters():
    if p.dim() > 1:
        nn.init.xavier_normal_(p)

optim = ScheduledOptim(
    Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-09),
    kwargs['d_model'], kwargs['n_warmup_steps'])

criterion = nn.CrossEntropyLoss(ignore_index=0)

# + id="d3FvWIMWpRpl" colab_type="code" colab={}
print_every = 500
model.train()

lowest_val = 1e9
train_losses = []
val_losses = []
total_step = 0

# + id="ND4J-jJbxjCa" colab_type="code" colab={}
# for epoch in range(kwargs['num_epochs']):
epoch = 1

# + id="EOkl_I_exku2" colab_type="code" colab={}
pbar = tqdm(total=print_every, leave=False)
total_loss = 0
train_loader.dataset.shuffle_batches()

# for step, (src, src_key_padding_mask, tgt, tgt_key_padding_mask) in enumerate(iter(train_loader)):
step, (src, src_key_padding_mask, tgt, tgt_key_padding_mask) = next(enumerate(iter(train_loader)))

# + id="Ze5Xan60zjN3" colab_type="code" colab={}
print(src.shape, tgt.shape)

total_step += 1

src, src_key_padding_mask = src[0].to('cuda'), src_key_padding_mask[0].to('cuda')
tgt, tgt_key_padding_mask = tgt[0].to('cuda'), tgt_key_padding_mask[0].to('cuda')

print(src.shape)
print(tgt.shape)

# + id="wUjqMCvt3U_J" colab_type="code" colab={}
memory_key_padding_mask = src_key_padding_mask.clone()

# + id="G5kF69V45F0X" colab_type="code" colab={}
tgt_inp, tgt_out = tgt[:, :-1], tgt[:, 1:] # shift for 1 position

# + id="mOf5mPfJ5HYK" colab_type="code" colab={}
tgt_inp.shape

# + id="5yJfGknU5Lh1" colab_type="code" colab={}
tgt_out.shape


# + id="E2aq9-DW57CB" colab_type="code" colab={}
def gen_nopeek_mask(length):
    mask = rearrange(torch.triu(torch.ones(length, length)) == 1, 'h w -> w h')
    mask = rearrange(torch.triu(torch.ones(length, length)) == 1, 'h w -> w h')
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask

gen_nopeek_mask(tgt_inp.shape[1])

# + id="E2MF0MHn8VKv" colab_type="code" colab={}
tgt_mask = gen_nopeek_mask(tgt_inp.shape[1]).to('cuda')

# + id="IFevNzG_yl_W" colab_type="code" colab={}
optim.zero_grad()
outputs = model(src, tgt_inp, src_key_padding_mask, tgt_key_padding_mask[:, :-1], memory_key_padding_mask, tgt_mask)
loss = criterion(rearrange(outputs, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))

# + id="J2mjXkUB8lxT" colab_type="code" colab={}
outputs.shape

# + id="K0Wo1KoS8d1G" colab_type="code" colab={}
rearrange(outputs, 'b t v -> (b t) v').shape

# + id="XA9eENjWI5c9" colab_type="code" colab={}
rearrange(tgt_out, 'b o -> (b o)').shape

# + id="aug6U-WOAzcR" colab_type="code" colab={}
rearrange(np.array([]))

# + id="1CFZCC2f83Nt" colab_type="code" colab={}


# + id="ZNdLwMFc8cjO" colab_type="code" colab={}
loss.backward()
optim.step_and_update_lr()

total_loss += loss.item()
train_losses.append((step, loss.item()))
pbar.update(1)
if step % print_every == print_every - 1:
    pbar.close()
    print(f'Epoch [{epoch + 1} / {num_epochs}] \t Step [{step + 1} / {len(train_loader)}] \t '
          f'Train Loss: {total_loss / print_every}')
    total_loss = 0

    pbar = tqdm(total=print_every, leave=False)

# + id="QhUT2foyvnFy" colab_type="code" colab={}
pbar.close()
val_loss = validate(valid_loader, model, criterion)
val_losses.append((total_step, val_loss))
if val_loss < lowest_val:
    lowest_val = val_loss
    torch.save(model, 'output/transformer.pth')
print(f'Val Loss: {val_loss}')

StarcoderdataJupyterScriptsDedupFiltered