opt = config['opt']

optimizer = config['optimizer']

scheduler = config['scheduler']

writer = config['writer']

scaler = config['scaler']

pad_label_id = config['pad_label_id']

criterion = nn.CrossEntropyLoss(ignore_index=pad_label_id).to(opt.device)

n_batches = len(train_loader)

prog = Progbar(target=n_batches)

# train one epoch

model.train()

train_loss = 0.

st_time = time.time()

optimizer.zero_grad()

for local_step, (x,y) in enumerate(train_loader):

global_step = (len(train_loader) * epoch_i) + local_step

x = to_device(x, opt.device)

y = to_device(y, opt.device)

if opt.use_crf:

with autocast(enabled=opt.use_amp):

if opt.use_profiler:

with profiler.profile(profile_memory=True, record_shapes=True) as prof:

logits, prediction = model(x)

print(prof.key_averages().table(sort_by="self_cpu_memory_usage", row_limit=10))

else:

logits, prediction = model(x)

mask = torch.sign(torch.abs(x[0])).to(torch.uint8).to(opt.device)

log_likelihood = model.crf(logits, y, mask=mask, reduction='mean')

loss = -1 * log_likelihood

if opt.gradient_accumulation_steps > 1:

loss = loss / opt.gradient_accumulation_steps

else:

with autocast(enabled=opt.use_amp):

if opt.use_profiler:

with profiler.profile(profile_memory=True, record_shapes=True) as prof:

logits = model(x)

print(prof.key_averages().table(sort_by="self_cpu_memory_usage", row_limit=10))

else:

logits = model(x)

# reshape for computing loss

logits_view = logits.view(-1, model.label_size)

y_view = y.view(-1)

loss = criterion(logits_view, y_view)

if opt.gradient_accumulation_steps > 1:

loss = loss / opt.gradient_accumulation_steps

# back-propagation - begin

scaler.scale(loss).backward()

if (local_step + 1) % opt.gradient_accumulation_steps == 0:

scaler.unscale_(optimizer)

torch.nn.utils.clip_grad_norm_(model.parameters(), opt.max_grad_norm)

scaler.step(optimizer)

scaler.update()

optimizer.zero_grad()

if opt.use_transformers_optimizer: scheduler.step()

# back-propagation - end

train_loss += loss.item()

if writer: writer.add_scalar('Loss/train', loss.item(), global_step)

curr_lr = scheduler.get_last_lr()[0] if scheduler else optimizer.param_groups[0]['lr']

prog.update(local_step+1,

[('global step', global_step),

('train curr loss', loss.item()),

('lr', curr_lr)])

train_loss = train_loss / n_batches

# evaluate

eval_ret = evaluate(model, config, val_loader)

eval_loss = eval_ret['loss']

eval_f1 = eval_ret['f1']

curr_time = time.time()

elapsed_time = (curr_time - st_time) / 60

st_time = curr_time

curr_lr = scheduler.get_last_lr()[0] if scheduler else optimizer.param_groups[0]['lr']

logger.info('{:3d} epoch | {:5d}/{:5d} | train loss : {:10.6f}, valid loss {:10.6f}, valid f1 {:.4f}| lr :{:7.6f} | {:5.2f} min elapsed'.\

format(epoch_i, local_step+1, len(train_loader), train_loss, eval_loss, eval_f1, curr_lr, elapsed_time))

if writer:

writer.add_scalar('Loss/valid', eval_loss, global_step)

writer.add_scalar('F1/valid', eval_f1, global_step)

writer.add_scalar('LearningRate/train', curr_lr, global_step)

model.eval()

opt = config['opt']

pad_label_id = config['pad_label_id']

eval_loss = 0.

criterion = nn.CrossEntropyLoss(ignore_index=pad_label_id).to(opt.device)

n_batches = len(val_loader)

prog = Progbar(target=n_batches)

preds = None

ys = None

with torch.no_grad():

for i, (x,y) in enumerate(val_loader):

x = to_device(x, opt.device)

y = to_device(y, opt.device)

if opt.use_crf:

logits, prediction = model(x)

mask = torch.sign(torch.abs(x[0])).to(torch.uint8).to(opt.device)

log_likelihood = model.crf(logits, y, mask=mask, reduction='mean')

loss = -1 * log_likelihood

else:

logits = model(x)

loss = criterion(logits.view(-1, model.label_size), y.view(-1))

if preds is None:

if opt.use_crf: preds = to_numpy(prediction)

else: preds = to_numpy(logits)

ys = to_numpy(y)

else:

if opt.use_crf: preds = np.append(preds, to_numpy(prediction), axis=0)

else: preds = np.append(preds, to_numpy(logits), axis=0)

ys = np.append(ys, to_numpy(y), axis=0)

eval_loss += loss.item()

prog.update(i+1,

[('eval curr loss', loss.item())])

eval_loss = eval_loss / n_batches

if not opt.use_crf: preds = np.argmax(preds, axis=2)

# compute measure using seqeval

labels = model.labels

ys_lbs = [[] for _ in range(ys.shape[0])]

preds_lbs = [[] for _ in range(ys.shape[0])]

for i in range(ys.shape[0]): # foreach sentence

for j in range(ys.shape[1]): # foreach token

if ys[i][j] != pad_label_id:

ys_lbs[i].append(labels[ys[i][j]])

preds_lbs[i].append(labels[preds[i][j]])

ret = {

"loss": eval_loss,

"precision": precision_score(ys_lbs, preds_lbs),

"recall": recall_score(ys_lbs, preds_lbs),

"f1": f1_score(ys_lbs, preds_lbs),

"report": classification_report(ys_lbs, preds_lbs, digits=4),

}

print(ret['report'])

opt = config['opt']

optimizer = config['optimizer']

checkpoint_path = opt.save_path

with open(checkpoint_path, 'wb') as f:

checkpoint = model.state_dict()

opt = config['opt']

if config['emb_class'] == 'glove':

opt.train_path = os.path.join(opt.data_dir, 'train.txt.ids')

opt.valid_path = os.path.join(opt.data_dir, 'valid.txt.ids')

if config['emb_class'] in ['bert', 'distilbert', 'albert', 'roberta', 'bart', 'electra']:

opt.train_path = os.path.join(opt.data_dir, 'train.txt.fs')

opt.valid_path = os.path.join(opt.data_dir, 'valid.txt.fs')

if config['emb_class'] == 'elmo':

opt.train_path = os.path.join(opt.data_dir, 'train.txt.ids')

opt.valid_path = os.path.join(opt.data_dir, 'valid.txt.ids')

opt.label_path = os.path.join(opt.data_dir, opt.label_filename)

opt.pos_path = os.path.join(opt.data_dir, opt.pos_filename)

opt = config['opt']

if config['emb_class'] == 'glove':

DatasetClass = CoNLLGloveDataset

if config['emb_class'] in ['bert', 'distilbert', 'albert', 'roberta', 'bart', 'electra']:

DatasetClass = CoNLLBertDataset

if config['emb_class'] == 'elmo':

DatasetClass = CoNLLElmoDataset

train_loader = prepare_dataset(config, opt.train_path, DatasetClass, sampling=True, num_workers=2, hp_search_bsz=hp_search_bsz)

valid_loader = prepare_dataset(config, opt.valid_path, DatasetClass, sampling=False, num_workers=2, batch_size=opt.eval_batch_size)

opt = config['opt']

embed_list = list(bert_model.embeddings.parameters())

# note that 'distilbert' has no encoder.layer, so don't use bert_remove_layers for distilbert.

layer_list = bert_model.encoder.layer

opt = config['opt']

remove_layers = opt.bert_remove_layers

# drop layers

if remove_layers is not "":

embed_list, layer_list = get_bert_embed_layer_list(config, bert_model)

layer_indexes = [int(x) for x in remove_layers.split(",")]

layer_indexes.sort(reverse=True)

for layer_idx in layer_indexes:

if layer_idx < 0 or layer_idx >= bert_config.num_hidden_layers: continue

del(layer_list[layer_idx])

logger.info("%s layer removed" % (layer_idx))

if len(layer_indexes) > 0:

opt = config['opt']

emb_non_trainable = not opt.embedding_trainable

if config['emb_class'] == 'glove':

if config['enc_class'] == 'bilstm':

model = GloveLSTMCRF(config, opt.embedding_path, opt.label_path, opt.pos_path,

emb_non_trainable=emb_non_trainable, use_crf=opt.use_crf, use_char_cnn=opt.use_char_cnn)

if config['enc_class'] == 'densenet':

model = GloveDensenetCRF(config, opt.embedding_path, opt.label_path, opt.pos_path,

emb_non_trainable=emb_non_trainable, use_crf=opt.use_crf, use_char_cnn=opt.use_char_cnn)

if config['emb_class'] in ['bert', 'distilbert', 'albert', 'roberta', 'bart', 'electra']:

from transformers import AutoTokenizer, AutoConfig, AutoModel

bert_tokenizer = AutoTokenizer.from_pretrained(opt.bert_model_name_or_path,

do_lower_case=opt.bert_do_lower_case)

bert_model = AutoModel.from_pretrained(opt.bert_model_name_or_path,

from_tf=bool(".ckpt" in opt.bert_model_name_or_path))

bert_config = bert_model.config

# bert model reduction

reduce_bert_model(config, bert_model, bert_config)

ModelClass = BertLSTMCRF

model = ModelClass(config, bert_config, bert_model, bert_tokenizer, opt.label_path, opt.pos_path,

use_crf=opt.use_crf, use_pos=opt.bert_use_pos, disable_lstm=opt.bert_disable_lstm, feature_based=opt.bert_use_feature_based)

if config['emb_class'] == 'elmo':

from allennlp.modules.elmo import Elmo

elmo_model = Elmo(opt.elmo_options_file, opt.elmo_weights_file, 2, dropout=0)

model = ElmoLSTMCRF(config, elmo_model, opt.embedding_path, opt.label_path, opt.pos_path,

emb_non_trainable=emb_non_trainable, use_crf=opt.use_crf, use_char_cnn=opt.use_char_cnn)

model.to(opt.device)

print(model)

logger.info("[model prepared]")

opt = config['opt']

lr = opt.lr

# for optuna

if hp_search_lr: lr = hp_search_lr

optimizer = torch.optim.AdamW(model.parameters(), lr=lr, eps=opt.adam_epsilon, weight_decay=opt.weight_decay)

scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=optimizer, gamma=opt.lr_decay_rate)

if opt.use_transformers_optimizer:

from transformers import AdamW, get_linear_schedule_with_warmup

num_training_steps_for_epoch = len(train_loader) // opt.gradient_accumulation_steps

num_training_steps = num_training_steps_for_epoch * opt.epoch

num_warmup_steps = num_training_steps_for_epoch * opt.warmup_epoch

logger.info("(num_training_steps_for_epoch, num_training_steps, num_warmup_steps): ({}, {}, {})".\

format(num_training_steps_for_epoch, num_training_steps, num_warmup_steps))

no_decay = ['bias', 'LayerNorm.weight']

optimizer_grouped_parameters = [

{'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],

'weight_decay': opt.weight_decay},

{'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}

]

optimizer = AdamW(optimizer_grouped_parameters, lr=lr, eps=opt.adam_epsilon)

scheduler = get_linear_schedule_with_warmup(optimizer,

num_warmup_steps=num_warmup_steps,

num_training_steps=num_training_steps)

try:

writer = SummaryWriter(log_dir=opt.log_dir)

except:

writer = None

scaler = GradScaler()

logger.info("[optimizer, scheduler, summary writer, scaler prepared]")

if torch.cuda.is_available():

logger.info("%s", torch.cuda.get_device_name(0))

# set etc

torch.autograd.set_detect_anomaly(True)

# set config

config = load_config(opt)

config['opt'] = opt

logger.info("%s", config)

# set path

set_path(config)

# prepare train, valid dataset

train_loader, valid_loader = prepare_datasets(config)

with temp_seed(opt.seed):

# prepare model

model = prepare_model(config)

# create optimizer, scheduler, summary writer, scaler

optimizer, scheduler, writer, scaler = prepare_osws(config, model, train_loader)

config['optimizer'] = optimizer

config['scheduler'] = scheduler

config['writer'] = writer

config['scaler'] = scaler

# training

early_stopping = EarlyStopping(logger, patience=opt.patience, measure='f1', verbose=1)

local_worse_steps = 0

prev_eval_f1 = -float('inf')

best_eval_f1 = -float('inf')

for epoch_i in range(opt.epoch):

epoch_st_time = time.time()

eval_loss, eval_f1 = train_epoch(model, config, train_loader, valid_loader, epoch_i)

# early stopping

if early_stopping.validate(eval_f1, measure='f1'): break

if eval_f1 > best_eval_f1:

best_eval_f1 = eval_f1

if opt.save_path:

logger.info("[Best model saved] : {:10.6f}".format(best_eval_f1))

save_model(config, model)

# save finetuned bert model/config/tokenizer

if config['emb_class'] in ['bert', 'distilbert', 'albert', 'roberta', 'bart', 'electra']:

if not os.path.exists(opt.bert_output_dir):

os.makedirs(opt.bert_output_dir)

model.bert_tokenizer.save_pretrained(opt.bert_output_dir)

model.bert_model.save_pretrained(opt.bert_output_dir)

early_stopping.reset(best_eval_f1)

early_stopping.status()

# begin: scheduling, apply rate decay at the measure(ex, loss) getting worse for the number of deacy epoch steps.

if prev_eval_f1 >= eval_f1:

local_worse_steps += 1

else:

local_worse_steps = 0

logger.info('Scheduler: local_worse_steps / opt.lr_decay_steps = %d / %d' % (local_worse_steps, opt.lr_decay_steps))

if not opt.use_transformers_optimizer and \

epoch_i > opt.warmup_epoch and \

(local_worse_steps >= opt.lr_decay_steps or early_stopping.step() > opt.lr_decay_steps):

scheduler.step()

local_worse_steps = 0

prev_eval_f1 = eval_f1

if torch.cuda.is_available():

logger.info("%s", torch.cuda.get_device_name(0))

global gopt

opt = gopt

# set config

config = load_config(opt)

config['opt'] = opt

logger.info("%s", config)

# set path

set_path(config)

# set search spaces

lr = trial.suggest_float('lr', 1e-5, 1e-3, log=True)

bsz = trial.suggest_categorical('batch_size', [32, 64, 128])

seed = trial.suggest_int('seed', 17, 42)

epochs = trial.suggest_int('epochs', 1, opt.epoch)

# prepare train, valid dataset

train_loader, valid_loader = prepare_datasets(config, hp_search_bsz=bsz)

with temp_seed(seed):

# prepare model

model = prepare_model(config)

# create optimizer, scheduler, summary writer, scaler

optimizer, scheduler, writer, scaler = prepare_osws(config, model, train_loader, hp_search_lr=lr)

config['optimizer'] = optimizer

config['scheduler'] = scheduler

config['writer'] = writer

config['scaler'] = scaler

early_stopping = EarlyStopping(logger, patience=opt.patience, measure='f1', verbose=1)

best_eval_f1 = -float('inf')

for epoch in range(epochs):

eval_loss, eval_f1 = train_epoch(model, config, train_loader, valid_loader, epoch)

# early stopping

if early_stopping.validate(eval_f1, measure='f1'): break

if eval_f1 > best_eval_f1:

best_eval_f1 = eval_f1

early_stopping.reset(best_eval_f1)

early_stopping.status()

trial.report(eval_f1, epoch)

if trial.should_prune():

raise optuna.TrialPruned()

parser = argparse.ArgumentParser()

parser.add_argument('--config', type=str, default='configs/config-glove.json')

parser.add_argument('--data_dir', type=str, default='data/conll2003')

parser.add_argument('--embedding_filename', type=str, default='embedding.npy')

parser.add_argument('--label_filename', type=str, default='label.txt')

parser.add_argument('--pos_filename', type=str, default='pos.txt')

parser.add_argument('--device', type=str, default='cuda')

parser.add_argument('--batch_size', type=int, default=32)

parser.add_argument('--eval_batch_size', type=int, default=128)

parser.add_argument('--epoch', type=int, default=30)

parser.add_argument('--lr', type=float, default=1e-3)

parser.add_argument('--lr_decay_rate', type=float, default=1.0, help="Disjoint with --use_transformers_optimizer")

parser.add_argument('--lr_decay_steps', type=float, default=2, help="Number of decay epoch steps to be paitent. disjoint with --use_transformers_optimizer")

parser.add_argument('--warmup_epoch', type=int, default=4, help="Number of warmup epoch steps")

parser.add_argument('--patience', default=7, type=int, help="Max number of epoch to be patient for early stopping.")

parser.add_argument('--adam_epsilon', type=float, default=1e-8)

parser.add_argument('--weight_decay', type=float, default=0.01)

parser.add_argument('--gradient_accumulation_steps', type=int, default=1,

help="Number of updates steps to accumulate before performing a backward/update pass.")

parser.add_argument('--max_grad_norm', default=1.0, type=float, help="Max gradient norm.")

parser.add_argument('--save_path', type=str, default='pytorch-model-glove.pt')

parser.add_argument('--log_dir', type=str, default='runs')

parser.add_argument('--seed', default=42, type=int)

parser.add_argument('--use_crf', action='store_true', help="Add CRF layer")

parser.add_argument('--embedding_trainable', action='store_true', help="Set word embedding(Glove) trainable")

parser.add_argument('--use_char_cnn', action='store_true', help="Add Character features")

parser.add_argument('--use_transformers_optimizer', action='store_true', help="Use transformers AdamW, get_linear_schedule_with_warmup.")

parser.add_argument('--use_amp', action='store_true', help="Use automatic mixed precision.")

parser.add_argument('--use_profiler', action='store_true', help="Use profiler.")

# for BERT

parser.add_argument('--bert_model_name_or_path', type=str, default='bert-base-uncased',

help="Path to pre-trained model or shortcut name(ex, bert-base-uncased)")

parser.add_argument('--bert_do_lower_case', action='store_true',

help="Set this flag if you are using an uncased model.")

parser.add_argument('--bert_output_dir', type=str, default='bert-checkpoint',

help="The output directory where the model predictions and checkpoints will be written.")

parser.add_argument('--bert_use_feature_based', action='store_true',

help="Use BERT as feature-based, default fine-tuning")

parser.add_argument('--bert_disable_lstm', action='store_true',

help="Disable lstm layer")

parser.add_argument('--bert_use_pos', action='store_true', help="Add Part-Of-Speech features")

parser.add_argument('--bert_remove_layers', type=str, default='',

help="Specify layer numbers to remove during finetuning e.g. 8,9,10,11 to remove last 4 layers from BERT base(12 layers)")

# for ELMo

parser.add_argument('--elmo_options_file', type=str, default='embeddings/elmo_2x4096_512_2048cnn_2xhighway_5.5B_options.json')

parser.add_argument('--elmo_weights_file', type=str, default='embeddings/elmo_2x4096_512_2048cnn_2xhighway_5.5B_weights.hdf5')

# for Optuna

parser.add_argument('--hp_search', action='store_true',

help="Set this flag to use hyper-parameter search.")

parser.add_argument('--hp_trials', default=24, type=int,

help="Number of trials for hyper-parameter search.")

opt = parser.parse_args()

if opt.hp_search:

global gopt

gopt = opt

study = optuna.create_study(direction='maximize')

study.optimize(hp_search, n_trials=opt.hp_trials)

df = study.trials_dataframe(attrs=('number', 'value', 'params', 'state'))

print(df)

logger.info("study.best_params : %s", study.best_params)

logger.info("study.best_value : %s", study.best_value)

logger.info("study.best_trial : %s", study.best_trial) # for all, study.trials

else: