import time

import os

import numpy as np

import random

import torch

import torch.nn as nn

import torch.onnx

from torchtext import data

from args import get_args

from model import SmPlusPlus

from utils.relevancy_metrics import get_map_mrr

from trec_dataset import TrecDataset

from wiki_dataset import WikiDataset

args = get_args()

config = args

torch.manual_seed(args.seed)

def set_vectors(field, vector_path):

return field

args = get_args()

config = args

torch.manual_seed(args.seed)

torch.backends.cudnn.deterministic = True

if not args.cuda:

args.gpu = -1

if torch.cuda.is_available() and args.cuda:

print("Note: You are using GPU for training")

torch.cuda.set_device(args.gpu)

torch.cuda.manual_seed(args.seed)

if torch.cuda.is_available() and not args.cuda:

print("You have Cuda but you're using CPU for training.")

np.random.seed(args.seed)

random.seed(args.seed)

QID = data.Field(sequential=False)

QUESTION = data.Field(batch_first=True)

ANSWER = data.Field(batch_first=True)

LABEL = data.Field(sequential=False)

EXTERNAL = data.Field(sequential=True, tensor_type=torch.FloatTensor, batch_first=True, use_vocab=False,

postprocessing=data.Pipeline(lambda arr, _, train: [float(y) for y in arr]))

if config.dataset == 'TREC':

train, dev, test = TrecDataset.splits(QID, QUESTION, ANSWER, EXTERNAL, LABEL)

elif config.dataset == 'wiki':

train, dev, test = WikiDataset.splits(QID, QUESTION, ANSWER, EXTERNAL, LABEL)

else:

print("Unsupported dataset")

exit()

QID.build_vocab(train, dev, test)

QUESTION.build_vocab(train, dev, test)

ANSWER.build_vocab(train, dev, test)

LABEL.build_vocab(train, dev, test)

QUESTION = set_vectors(QUESTION, args.vector_cache)

ANSWER = set_vectors(ANSWER, args.vector_cache)

train_iter = data.Iterator(train, batch_size=args.batch_size, device=args.gpu, train=True, repeat=False,

sort=False, shuffle=True)

dev_iter = data.Iterator(dev, batch_size=args.batch_size, device=args.gpu, train=False, repeat=False,

sort=False, shuffle=False)

test_iter = data.Iterator(test, batch_size=args.batch_size, device=args.gpu, train=False, repeat=False,

sort=False, shuffle=False)

config.target_class = len(LABEL.vocab)

config.questions_num = len(QUESTION.vocab)

config.answers_num = len(ANSWER.vocab)

print("Dataset {} Mode {}".format(args.dataset, args.mode))

print("VOCAB num", len(QUESTION.vocab))

print("LABEL.target_class:", len(LABEL.vocab))

print("LABELS:", LABEL.vocab.itos)

print("Train instance", len(train))

print("Dev instance", len(dev))

print("Test instance", len(test))

if args.resume_snapshot:

if args.cuda:

model = torch.load(args.resume_snapshot, map_location=lambda storage, location: storage.cuda(args.gpu))

else:

model = torch.load(args.resume_snapshot, map_location=lambda storage, location: storage)

else:

model = SmPlusPlus(config)

model.static_question_embed.weight.data.copy_(QUESTION.vocab.vectors)

model.nonstatic_question_embed.weight.data.copy_(QUESTION.vocab.vectors)

model.static_answer_embed.weight.data.copy_(ANSWER.vocab.vectors)

model.nonstatic_answer_embed.weight.data.copy_(ANSWER.vocab.vectors)

if args.cuda:

model.cuda()

print("Shift model to GPU")

parameter = filter(lambda p: p.requires_grad, model.parameters())

optimizer = torch.optim.Adadelta(parameter, lr=args.lr, weight_decay=args.weight_decay)

criterion = nn.CrossEntropyLoss()

early_stop = False

best_dev_map = 0

iterations = 0

iters_not_improved = 0

epoch = 0

start = time.time()

header = ' Time Epoch Iteration Progress (%Epoch) Loss Dev/Loss Accuracy Dev/Accuracy'

dev_log_template = ' '.join('{:>6.0f},{:>5.0f},{:>9.0f},{:>5.0f}/{:<5.0f} {:>7.0f}%,{:>8.6f},{:8.6f},{:12.4f},{:12.4f}'.split(','))

log_template = ' '.join('{:>6.0f},{:>5.0f},{:>9.0f},{:>5.0f}/{:<5.0f} {:>7.0f}%,{:>8.6f},{},{:12.4f},{}'.split(','))

os.makedirs(args.save_path, exist_ok=True)

os.makedirs(os.path.join(args.save_path, args.dataset), exist_ok=True)

print(header)

index2label = np.array(LABEL.vocab.itos)

index2qid = np.array(QID.vocab.itos)

index2question = np.array(ANSWER.vocab.itos)

while True:

if early_stop:

print("Early Stopping. Epoch: {}, Best Dev Acc: {}".format(epoch, best_dev_map))

break

epoch += 1

train_iter.init_epoch()

n_correct, n_total = 0, 0

for batch_idx, batch in enumerate(train_iter):

iterations += 1

model.train(); optimizer.zero_grad()

scores = model(batch.question, batch.answer, batch.ext_feat)

n_correct += (torch.max(scores, 1)[1].view(batch.label.size()).data == batch.label.data).sum()

n_total += batch.batch_size

train_acc = 100. * n_correct / n_total

loss = criterion(scores, batch.label)

loss.backward()

optimizer.step()

# Evaluate performance on validation set

if iterations % args.dev_every == 1:

# switch model into evaluation mode

model.eval()

dev_iter.init_epoch()

n_dev_correct = 0

dev_losses = []

qids = []

predictions = []

labels = []

for dev_batch_idx, dev_batch in enumerate(dev_iter):

qid_array = index2qid[np.transpose(dev_batch.qid.cpu().data.numpy())]

true_label_array = index2label[np.transpose(dev_batch.label.cpu().data.numpy())]

scores = model(dev_batch.question, dev_batch.answer, dev_batch.ext_feat)

n_dev_correct += (torch.max(scores, 1)[1].view(dev_batch.label.size()).data == dev_batch.label.data).sum()

dev_loss = criterion(scores, dev_batch.label)

dev_losses.append(dev_loss.data[0])

index_label = np.transpose(torch.max(scores, 1)[1].view(dev_batch.label.size()).cpu().data.numpy())

label_array = index2label[index_label]

# get the relevance scores

score_array = scores[:, 2].cpu().data.numpy()

qids.extend(qid_array.tolist())

predictions.extend(score_array.tolist())

labels.extend(true_label_array.tolist())

dev_map, dev_mrr = get_map_mrr(qids, predictions, labels)

print(dev_log_template.format(time.time() - start,

epoch, iterations, 1 + batch_idx, len(train_iter),

100. * (1 + batch_idx) / len(train_iter), loss.data[0],

sum(dev_losses) / len(dev_losses), train_acc, dev_map))

# Update validation results

if dev_map > best_dev_map:

iters_not_improved = 0

best_dev_map = dev_map

snapshot_path = os.path.join(args.save_path, args.dataset, args.mode+'_best_model.pt')

torch.save(model, snapshot_path)

else:

iters_not_improved += 1

if iters_not_improved >= args.patience:

early_stop = True

break

if iterations % args.log_every == 1:

# print progress message

print(log_template.format(time.time() - start,

epoch, iterations, 1 + batch_idx, len(train_iter),

100. * (1 + batch_idx) / len(train_iter), loss.data[0], ' ' * 8,

n_correct / n_total * 100, ' ' * 12))