import pickle

import json

import os

import math

import torch

import numpy as np

import matplotlib.pyplot as plt

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

from torch.hub import download_url_to_file

import torch.utils.data

# pip install nltk

import nltk

from nltk.tokenize import sent_tokenize, word_tokenize, RegexpTokenizer

from nltk.corpus import stopwords

# nltk.download('stopwords')

stop_words = stopwords.words('english')

# print(stop_words)

# nltk.download() #downlaod punkt manualy

BATCH_SIZE = 64

EPOCHS = 500

LEARNING_RATE = 1e-3

RNN_INPUT_SIZE = 256

RNN_HIDDEN_SIZE = 256 # 512

RNN_LAYERS = 1

RNN_DROPOUT = 0.3

PACKING = False # True

run_path = ''

DEVICE = 'cpu'

if torch.cuda.is_available():

DEVICE = 'cuda'

MIN_SENTENCE_LEN = 3

MAX_SENTENCE_LEN = 20

MAX_LEN = 200 # 200 # limit max number of samples otherwise too slow training (on GPU use all samples / for final training)

if DEVICE == 'cuda':

MAX_LEN = 10000

PATH_DATA = './data'

os.makedirs('./results', exist_ok=True)

os.makedirs(PATH_DATA, exist_ok=True)

# class DatasetCustom(torch.utils.data.Dataset):

# def __init__(self):

# # if not os.path.exists(f'{PATH_DATA}/recipes_raw_nosource_epi.json'):

# # download_url_to_file('https://www.yellowrobot.xyz/share/recipes_raw_nosource_epi.json', progress=True)

# # doesn't work

# # download_url_to_file('https://www.kaggle.com/akmittal/quotes-dataset/download',

# # dst=f'{PATH_DATA}/a.zip', progress=True)

# # with open(f'{PATH_DATA}/recipes_raw_nosource_epi.json') as fp:

# # data_json = json.load(fp)

# with open(f'{PATH_DATA}/quotes.json', encoding="utf-8") as fp:

# data_json = json.load(fp)

#

# self.sentences = []

# self.lengths = [] # lengths of all sentences

# self.words_to_idxes = {}

# self.words_counts = {}

# self.idxes_to_words = {}

#

# for quote_obj in data_json:

# quote = quote_obj['Quote']

# sentences = sent_tokenize(quote)

# for sentence in sentences:

# # split sentence into words and make them lowercase

# words = str.split(sentence.lower())

# # remove punctuation only at the end of word - e.x. "don't'" -> "don't"

# words = [w[:-1] if str.isalpha(w[-1]) is False else w for w in words]

# # filter out stop words

# words = [w for w in words if w not in stop_words]

# # filter one character tokens

# words = [w for w in words if len(w) > 1]

# # remove bad words i'm, i've, .., 10,00 (just didn't overcomplicate with other processing)

# words = [w for w in words if w not in ["i'm", "i've", "..", "10,00"]]

#

# if len(words) > MAX_SENTENCE_LEN:

# words = words[:MAX_SENTENCE_LEN]

# if len(words) < MIN_SENTENCE_LEN:

# continue

# sentence_tokens = []

# for word in words:

# if word not in self.words_to_idxes:

# self.words_to_idxes[word] = len(self.words_to_idxes)

# self.idxes_to_words[self.words_to_idxes[word]] = word

# self.words_counts[word] = 0

# self.words_counts[word] += 1

# sentence_tokens.append(self.words_to_idxes[word])

# self.sentences.append(sentence_tokens)

# self.lengths.append(len(sentence_tokens))

# if MAX_LEN is not None and len(self.sentences) > MAX_LEN:

# break

#

# # after checking freq dist removed characters

# words_freq = dict(sorted(self.words_counts.items(), key=lambda item: item[1]))

# n = 40

# words = list(words_freq.keys())[-n:]

# vals = list(words_freq.values())[-n:]

# plt.barh(words, vals)

# plt.show()

# # sum(vals) = 375 now, before it was 1122

#

# self.max_length = np.max(self.lengths) + 1 # longest sentence length + 1

# self.end_token = '[END]'

# self.words_to_idxes[self.end_token] = len(self.words_to_idxes)

# self.idxes_to_words[self.words_to_idxes[self.end_token]] = self.end_token

# self.words_counts[self.end_token] = len(self.sentences)

#

# self.max_classes_tokens = len(self.words_to_idxes) # unique words amount

#

# word_counts = np.array(list(self.words_counts.values()))

# self.weights = (1.0 / word_counts) * np.sum(word_counts) * 0.5 # more frequent word == smaller weight

#

# print(f'self.sentences: {len(self.sentences)}')

# print(f'self.max_length: {self.max_length}')

# print(f'self.max_classes_tokens: {self.max_classes_tokens}')

#

# print('Example sentences:')

# samples = np.random.choice(self.sentences, 5)

# for each in samples:

# print(' '.join([self.idxes_to_words[it] for it in each]))

#

# def __len__(self):

# return len(self.sentences)

#

# def __getitem__(self, idx):

# # data - abcdef

# # x - abcde<end>

# # y - bcdef<end>

#

# np_x_idxes = np.array(self.sentences[idx][:-1] + [self.words_to_idxes[self.end_token]])

# np_x_padded = np.zeros((self.max_length, self.max_classes_tokens))

# np_x_padded[np.arange(len(np_x_idxes)), np_x_idxes] = 1.0

#

# np_y_idxes = np.array(self.sentences[idx][1:] + [self.words_to_idxes[self.end_token]])

# np_y_padded = np.zeros((self.max_length, self.max_classes_tokens))

# np_y_padded[np.arange(len(np_y_idxes)), np_y_idxes] = 1.0

#

# np_length = self.lengths[idx]

#

# # print([self.idxes_to_words[idx] for idx in np_x_idxes])

# # print([self.idxes_to_words[idx] for idx in np_y_idxes])

# return np_x_padded, np_y_padded, np_length

# # zero seed will result on having always same random split

# torch.manual_seed(0)

# dataset_full = DatasetCustom()

# dataset_train, dataset_test = torch.utils.data.random_split(

# dataset_full, lengths=[int(len(dataset_full) * 0.8), len(dataset_full) - int(len(dataset_full) * 0.8)])

# torch.seed()

#

# data_loader_train = torch.utils.data.DataLoader(

# dataset=dataset_train,

# batch_size=BATCH_SIZE,

# shuffle=True

# )

# data_loader_test = torch.utils.data.DataLoader(

# dataset=dataset_test,

# batch_size=BATCH_SIZE,

# shuffle=False

# )

# TAKEN FROM TRANSFORMER.PY

class DatasetCustom(torch.utils.data.Dataset):

return np_x_padded, np_y_padded, np_length

if not os.path.exists(f'{PATH_DATA}/dataset_full.pt'):

download_url_to_file('http://www.yellowrobot.xyz/share/dataset_full.pt', f'{PATH_DATA}/dataset_full.pt',

progress=True)

with open(f'{PATH_DATA}/dataset_full.pt', 'rb') as fp:

dataset_full = pickle.load(fp)

torch.manual_seed(0)

dataset_train, dataset_test = torch.utils.data.random_split(

dataset_full, lengths=[int(len(dataset_full) * 0.8), len(dataset_full) - int(len(dataset_full) * 0.8)])

torch.seed()

data_loader_train = torch.utils.data.DataLoader(

)

data_loader_test = torch.utils.data.DataLoader(

)

########################################################################################

class GRUCell(torch.nn.Module):

return t_h_packed

class RNNCell(torch.nn.Module):

return t_h_packed

class Model(torch.nn.Module):

return y_prim_packed

model = Model()

model = model.to(DEVICE)

optimizer = torch.optim.RMSprop(model.parameters(), lr=1e-4)

metrics = {}

best_test_loss = float('Inf')

for stage in ['train', 'test']:

for metric in [

'loss',

'acc'

]:

metrics[f'{stage}_{metric}'] = []

for epoch in range(1, EPOCHS + 1):

for data_loader in [data_loader_train, data_loader_test]:

metrics_epoch = {key: [] for key in metrics.keys()}

stage = 'train'

if data_loader == data_loader_test:

stage = 'test'

for x, y, lengths in data_loader:

x = x.float().to(DEVICE)

y = y.float().to(DEVICE)

idxes = torch.argsort(lengths, descending=True)

lengths = lengths[idxes]

max_len = int(lengths.max())

# sort sentences by length desc and slice

# in x last word is either empty or 'END', and in y it is shifted first word)

x = x[idxes, :max_len]

y = y[idxes, :max_len]

x_packed = pack_padded_sequence(x, lengths, batch_first=True)

y_packed = pack_padded_sequence(y, lengths, batch_first=True)

y_prim_packed = model.forward(x_packed)

weights = torch.from_numpy(dataset_full.weights[torch.argmax(y_packed.data, dim=1).cpu().numpy()])

weights = weights.unsqueeze(dim=1).to(DEVICE)

loss = -torch.mean(weights * y_packed.data * torch.log(y_prim_packed.data + 1e-8))

metrics_epoch[f'{stage}_loss'].append(loss.item()) # Tensor(0.1) => 0.1f

if data_loader == data_loader_train:

loss.backward()

optimizer.step()

optimizer.zero_grad()

np_y_prim = y_prim_packed.data.cpu().data.numpy()

np_y = y_packed.data.cpu().data.numpy()

idx_y = np.argmax(np_y, axis=1)

idx_y_prim = np.argmax(np_y_prim, axis=1)

acc = np.average((idx_y == idx_y_prim) * 1.0)

metrics_epoch[f'{stage}_acc'].append(acc)

metrics_strs = []

for key in metrics_epoch.keys():

if stage in key:

value = np.mean(metrics_epoch[key])

metrics[key].append(value)

metrics_strs.append(f'{key}: {round(value, 2)}')

print(f'epoch: {epoch} {" ".join(metrics_strs)}')

# validation

if data_loader == data_loader_test:

y_prim_unpacked, lengths_unpacked = pad_packed_sequence(y_prim_packed.cpu(), batch_first=True)

y_prim_unpacked = y_prim_unpacked[0]

x = x[0]

y = y[0]

y_prim_idxes = np.argmax(y_prim_unpacked[:lengths_unpacked[0]].cpu().data.numpy(), axis=1).tolist()

x_idxes = np.argmax(x[:lengths_unpacked[0]].cpu().data.numpy(), axis=1).tolist()

y_idxes = np.argmax(y[:lengths_unpacked[0]].cpu().data.numpy(), axis=1).tolist()

print('Validation:')

print('x: ' + ' '.join([dataset_full.idxes_to_words[it] for it in x_idxes]))

print('y: ' + ' '.join([dataset_full.idxes_to_words[it] for it in y_idxes]))

print('y_prim: ' + ' '.join([dataset_full.idxes_to_words[it] for it in y_prim_idxes]))

print(' ')

if best_test_loss > loss.item():

best_test_loss = loss.item()

torch.save(model.cpu().state_dict(), f'./results/model-{epoch}.pt')

model = model.to(DEVICE)

plt.figure(figsize=(12, 5))

plts = []

c = 0

for key, value in metrics.items():

plts += plt.plot(value, f'C{c}', label=key)

ax = plt.twinx()

c += 1

plt.legend(plts, [it.get_label() for it in plts])

plt.savefig(f'./results/epoch-{epoch}.png')

plt.show()

"""

# # filter punctions

# words = RegexpTokenizer(r'\w+').tokenize(sentence.lower())

# e.x. 'Don\'t cry because it\'s over, smile because it happened.'

# 1. RegexpTokenizer(r'\w+').tokenize(s)

# ['Don', 't', 'cry', 'because', 'it', 's', 'over', 'smile', 'because', 'it', 'happened']

# 2. words = word_tokenize(sentence.lower()) -> words=[word.lower() for word in words if word.isalpha()]

# ['do', "n't", 'cry', 'because', 'it', "'s", 'over', ',', 'smile', 'because', 'it', 'happened', '.']

# ['do', 'cry', 'because', 'it', 'over', 'smile', 'because', 'it', 'happened']

"""