from contextlib import ExitStack

from pathlib import Path

from typing import Tuple

import numpy as np

import torch

import torch.multiprocessing

import torch.nn as nn

from labml import experiment, tracker

from numpy import ndarray

from torch.utils.data import DataLoader

from tqdm import tqdm

from neuro_comma import augmentation

from neuro_comma.argparser import parse_train_arguments

from neuro_comma.dataset import RepunctDataset

from neuro_comma.logger import (log_args, log_target_test_metrics,

log_test_metrics, log_text, log_train_epoch,

log_val_epoch)

from neuro_comma.model import CorrectionModel

from neuro_comma.pretrained import PRETRAINED_MODELS

from neuro_comma.utils import (export_params, get_last_epoch_params,

get_last_pretrained_weight_path,

get_model_save_path, load_params, save_weights)

# https://github.com/pytorch/pytorch/issues/11201

torch.multiprocessing.set_sharing_strategy('file_system')

args = parse_train_arguments()

# for reproducibility

if args.seed:

torch.manual_seed(args.seed)

torch.backends.cudnn.deterministic = True

torch.backends.cudnn.benchmark = False

np.random.seed(args.seed)

models_root = Path(args.save_dir)

model_save_path = get_model_save_path(models_root, args)

# MODEL

print('Loading model...')

TARGETS = args.targets

DEVICE = torch.device('cuda' if (args.cuda and torch.cuda.is_available()) else 'cpu')

MODEL_SAVE_NAME = model_save_path.stem

if not (args.resume or args.fine_tune):

MODEL = CorrectionModel(args.pretrained_model,

targets=TARGETS,

freeze_pretrained=args.freeze_pretrained,

lstm_dim=args.lstm_dim)

else:

orig_model_dir = models_root / args.model_name

orig_params = load_params(orig_model_dir)

MODEL = CorrectionModel(pretrained_model=orig_params['pretrained_model'],

targets=orig_params['targets'],

freeze_pretrained=args.freeze_pretrained,

lstm_dim=orig_params['lstm_dim'])

pretrained_weights = get_last_pretrained_weight_path(orig_model_dir)

MODEL.load(pretrained_weights)

if args.fine_tune and (len(orig_params['targets']) != len(TARGETS)):

MODEL.modify_last_linear(in_features=MODEL.hidden_size * 2,

out_features=len(TARGETS))

MODEL.to(DEVICE)

WEIGHTS = torch.FloatTensor(args.weights).to(DEVICE) if args.weights else None

CRITERION = nn.CrossEntropyLoss(weight=WEIGHTS)

OPTIMIZER = torch.optim.Adam(MODEL.parameters(), lr=args.lr, weight_decay=args.decay)

print('Model was loaded.')

# TOKENIZER

print('Loading tokenizer...')

tokenizer = PRETRAINED_MODELS[args.pretrained_model][1].from_pretrained(args.pretrained_model)

token_style = PRETRAINED_MODELS[args.pretrained_model][3]

SEQ_LEN = args.sequence_length

print('Tokenizer was loaded.')

# CONFIG AUGMENTATION

AUG_RATE = args.augment_rate

AUG_TYPE = args.augment_type

augmentation.tokenizer = tokenizer

augmentation.sub_style = args.sub_style

augmentation.alpha_sub = args.alpha_sub

augmentation.alpha_del = args.alpha_del

# DATASETS

data_loader_params = {

'batch_size': args.batch_size,

'shuffle': True,

'num_workers': 1

}

print('Loading train data...')

train_dataset = RepunctDataset(args.train_data, tokenizer=tokenizer, targets=TARGETS,

sequence_len=SEQ_LEN, token_style=token_style,

is_train=True, augment_rate=AUG_RATE,

augment_type=AUG_TYPE, debug=True)

train_loader = DataLoader(train_dataset, **data_loader_params)

print('Loading validation data...')

val_dataset = RepunctDataset(args.val_data, tokenizer=tokenizer, targets=TARGETS,

sequence_len=SEQ_LEN, token_style=token_style,

is_train=True, debug=True)

val_loader = DataLoader(val_dataset, **data_loader_params)

if args.test_data:

print('Loading test data...')

if args.test_data == args.val_data:

test_dataset = val_dataset

test_loader = val_loader

else:

test_dataset = RepunctDataset(args.test_data, tokenizer=tokenizer, targets=TARGETS,

sequence_len=SEQ_LEN, token_style=token_style,

is_train=True, augment_rate=AUG_RATE,

augment_type=AUG_TYPE)

test_loader = DataLoader(test_dataset, **data_loader_params)

print('Data was loaded.')

LOG_PATH = model_save_path / 'logs' / f"{args.model_name}_logs.txt"

LOG_PATH.parent.mkdir(parents=True, exist_ok=True)

WEIGHTS_SAVE_DIR = model_save_path / 'weights'

WEIGHTS_SAVE_DIR.mkdir(parents=True, exist_ok=True)

def train_epoch(model: nn.Module,

loader: DataLoader,

criterion: nn.Module,

optimizer: torch.optim.Optimizer

) -> Tuple[float, float]:

"""

Train single epoch

Args:

model (nn.Module): model instance

loader (DataLoader): data loader

criterion (nn.Module): criterion

optimizer (torch.optim.Optimizer): optimizer

Returns:

tuple[float, float]: train_loss, train_accuracy

"""

train_loss = 0.0

train_iteration = 0

correct = 0.

total = 0.

model.train()

for x, y, att, y_mask in tqdm(loader, desc='train'):

x = x.to(DEVICE)

y = y.view(-1).to(DEVICE)

att = att.to(DEVICE)

y_mask = y_mask.view(-1).to(DEVICE)

y_predict = model(x, att)

y_predict = y_predict.view(-1, y_predict.shape[2])

loss = criterion(y_predict, y)

y_predict = torch.argmax(y_predict, dim=1).view(-1)

correct += torch.sum(y_mask * (y_predict == y).long()).item()

optimizer.zero_grad()

train_loss += loss.item()

train_iteration += 1

loss.backward()

if args.gradient_clip > 0:

torch.nn.utils.clip_grad_norm_(model.parameters(), args.gradient_clip)

optimizer.step()

total += torch.sum(y_mask.view(-1)).item()

train_loss /= train_iteration

train_accuracy = correct / total

return train_loss, train_accuracy

def validate_epoch(model: nn.Module,

loader: DataLoader,

criterion: nn.Module) -> Tuple[float, float, float, float, float]:

"""

Validate sinlge epoch

Args:

model (nn.Module): model instance

loader (DataLoader): data loader (should be different from train loader)

criterion (nn.Module): criterion

Returns:

tuple[float, float, float, float, float]: validation_loss, validation_accuracy, f1, precision, recall

"""

num_iteration = 0

correct = 0.

total = 0.

val_loss = 0.0

# +1 for overall result

tp = np.zeros(1 + len(TARGETS), dtype=np.int64)

fp = np.zeros(1 + len(TARGETS), dtype=np.int64)

fn = np.zeros(1 + len(TARGETS), dtype=np.int64)

cm = np.zeros((len(TARGETS), len(TARGETS)), dtype=np.int64)

model.eval()

with torch.no_grad():

for x, y, att, y_mask in tqdm(loader, desc='eval'):

x = x.to(DEVICE)

y = y.view(-1).to(DEVICE)

att = att.to(DEVICE)

y_mask = y_mask.view(-1).to(DEVICE)

y_predict = model(x, att)

y_predict = y_predict.view(-1, y_predict.shape[2])

loss = criterion(y_predict, y)

val_loss += loss.item()

y_predict = torch.argmax(y_predict, dim=1).view(-1)

num_iteration += 1

y_mask = y_mask.view(-1)

correct += torch.sum(y_mask * (y_predict == y).long()).item()

total += torch.sum(y_mask).item()

for i in range(y.shape[0]):

if y_mask[i] == 0:

# we can ignore this because we know there won't be

# any punctuation in this position since we created

# this position due to padding or sub-word tokenization

continue

cor = y[i]

prd = y_predict[i]

if cor == prd:

tp[cor] += 1

else:

fn[cor] += 1

fp[prd] += 1

cm[cor][prd] += 1

# ignore first index which is for no punctuation

tp[-1] = np.sum(tp[1:])

fp[-1] = np.sum(fp[1:])

fn[-1] = np.sum(fn[1:])

precision = tp / (tp + fp)

recall = tp / (tp + fn)

f1 = 2 * precision * recall / (precision + recall)

val_loss = val_loss / num_iteration

val_acc = correct / total

log_test_metrics(LOG_PATH, precision, recall, f1, val_acc, cm)

non_O_keys = " + ".join(list(TARGETS)[1:])

targets = list(TARGETS) + [non_O_keys]

for i, target in enumerate(targets):

log_target_test_metrics(LOG_PATH, target, precision[i], recall[i], f1[i])

return val_loss, val_acc, f1, precision, recall

def calc_accuracy_metrics(model: nn.Module,

loader: DataLoader

) -> Tuple[ndarray, ndarray, ndarray, float, ndarray]:

"""

Calculate different accuracy metrics

Args:

model (nn.Module): model instance

loader (DataLoader): data loader

Returns:

tuple[ndarray, ndarray, ndarray, float, ndarray]: precision, recall, \

f1_score, accuracy, confusion_matrx

"""

num_iteration = 0

correct = 0

total = 0

model.eval()

# +1 for overall result

tp = np.zeros(1 + len(TARGETS), dtype=np.int64)

fp = np.zeros(1 + len(TARGETS), dtype=np.int64)

fn = np.zeros(1 + len(TARGETS), dtype=np.int64)

cm = np.zeros((len(TARGETS), len(TARGETS)), dtype=np.int64)

with torch.no_grad():

for x, y, att, y_mask in tqdm(loader, desc='test'):

x = x.to(DEVICE)

y = y.view(-1).to(DEVICE)

att = att.to(DEVICE)

y_mask = y_mask.view(-1).to(DEVICE)

y_predict = model(x, att)

y_predict = y_predict.view(-1, y_predict.shape[2])

y_predict = torch.argmax(y_predict, dim=1).view(-1)

num_iteration += 1

y_mask = y_mask.view(-1)

correct += int(torch.sum(y_mask * (y_predict == y).long()).item())

total += int(torch.sum(y_mask).item())

for i in range(y.shape[0]):

if y_mask[i] == 0:

# we can ignore this because we know there won't be

# any punctuation in this position since we created

# this position due to padding or sub-word tokenization

continue

cor = y[i]

prd = y_predict[i]

if cor == prd:

tp[cor] += 1

else:

fn[cor] += 1

fp[prd] += 1

cm[cor][prd] += 1

# ignore first index which is for no punctuation

tp[-1] = np.sum(tp[1:])

fp[-1] = np.sum(fp[1:])

fn[-1] = np.sum(fn[1:])

precision = tp / (tp + fp)

recall = tp / (tp + fn)

f1 = 2 * precision * recall / (precision + recall)

accuracy = correct / total

return precision, recall, f1, accuracy, cm

def train() -> None:

"""Train: global function"""

if not args.resume:

best_val_acc = 0.0

epochs = range(args.epoch)

else:

last_epoch, best_val_acc = get_last_epoch_params(orig_model_dir / 'weights')

epochs = range(last_epoch + 1, last_epoch + 1 + args.epoch)

# TODO: continue logging, remove print

with experiment.record(name=MODEL_SAVE_NAME, exp_conf=args.__dict__) if args.labml else ExitStack():

for epoch in epochs:

train_loss, train_acc = train_epoch(MODEL, train_loader, CRITERION, OPTIMIZER)

log_train_epoch(LOG_PATH, epoch, train_loss, train_acc)

val_loss, val_acc, f1, precision, recall = validate_epoch(MODEL, val_loader, CRITERION)

log_val_epoch(LOG_PATH, epoch, val_loss, val_acc)

if args.labml:

tracker.save(epoch, {'train_loss': train_loss,

'train_accuracy': train_acc,

'val_loss': val_loss,

'val_accuracy': val_acc,

'f1': f1,

'precision': precision,

'recall': recall})

if args.store_every_weight:

save_weights(MODEL, WEIGHTS_SAVE_DIR, epoch, val_acc)

elif args.store_best_weights and (val_acc > best_val_acc):

best_val_acc = val_acc

save_weights(MODEL, WEIGHTS_SAVE_DIR, epoch, val_acc)

log_text(LOG_PATH, f"Best validation Acc: {best_val_acc}")

def test() -> None:

"""Test: global function"""

# precision, recall, f1, accuracy, cm = calc_accuracy_metrics(MODEL, test_loader)

precision, recall, f1, accuracy, cm = calc_accuracy_metrics(MODEL, val_loader)

log_test_metrics(LOG_PATH, precision, recall, f1, accuracy, cm)

non_O_keys = " + ".join(list(TARGETS)[1:])

targets = list(TARGETS) + [non_O_keys]

for i, target in enumerate(targets):

log_target_test_metrics(LOG_PATH, target, precision[i], recall[i], f1[i])

if __name__ == '__main__':

export_params(args, model_save_path)

log_args(LOG_PATH, args)

train()

if args.test_data:

test()