""" Average a scalar over the nodes if we are in distributed training. We use this for distributed evaluation. """

if config.local_rank == -1:

return scalar

scalar_t = torch.tensor(scalar, dtype=torch.float, device=config.device) / torch.distributed.get_world_size()

torch.distributed.all_reduce(scalar_t, op=torch.distributed.ReduceOp.SUM)

""" Pad the dataset. This could be optimized by defining a Dataset class and padd only batches but this is simpler. """

max_l = max(len(x) for x in dataset["input_ids"])

for name in PADDED_INPUTS:

if name == "lm_labels":

dataset[name] = [x + [-1] * (max_l - len(x)) for x in dataset[name]]

elif name == "input_mask":

dataset[name] = [x + [0] * (max_l - len(x)) for x in dataset[name]]

else:

dataset[name] = [x + [padding] * (max_l - len(x)) for x in dataset[name]]

""" Build a sequence of input from 3 segments: persona, history and last reply """

bos, eos, speaker1, speaker2 = tokenizer.convert_tokens_to_ids(SPECIAL_TOKENS[:-1])

# todo rooh, in bert embedding we have separate embeddings for token_type_ids, input_ids and

# position so there shouldn't be any problem

speaker1 = 0

speaker2 = 1

instance = {}

sequence = [[bos] + list(chain(*persona))] + history + [

reply + ([eos] if with_eos else [])] # seq = [personas, history, reply] concatenate all persona sentences

sequence = [sequence[0]] + [[speaker2 if (len(sequence) - i) % 2 else speaker1] + s for i, s in

enumerate(sequence[1:])]

instance["input_ids"] = list(chain(*sequence))

instance["token_type_ids"] = [speaker2 if i % 2 else speaker1 for i, s in enumerate(sequence) for _ in

s] # the last for is for repeating the speaker1 and speaker2 for all tokens

instance["mc_token_ids"] = len(instance["input_ids"]) - 1

instance["lm_labels"] = [-1] * len(instance["input_ids"])

instance["input_mask"] = [1] * len(instance["input_ids"])

if lm_labels:

instance["lm_labels"] = ([-1] * sum(len(s) for s in sequence[:-1])) + [-1] + sequence[-1][

1:] # all -1 except for reply, reply is just the ids

""" Prepare the dataset for training and evaluation """

personachat = get_dataset(tokenizer, config.dataset_path, config.dataset_cache)

# todo remove soon

personachat["train"] = personachat["train"][:1000]

personachat["valid"] = personachat["valid"][:100]

logger.info("Build inputs and labels")

datasets = {"train": defaultdict(list), "valid": defaultdict(list)}

for dataset_name, dataset in personachat.items():

num_candidates = len(dataset[0]["utterances"][0]["candidates"])

if config.num_candidates > 0 and dataset_name == 'train':

num_candidates = min(config.num_candidates,

num_candidates) # for training we use 2 candidates but for validation 20

for dialog in dataset:

persona = dialog["personality"].copy()

for _ in range(config.personality_permutations):

for utterance in dialog["utterances"]:

history = utterance["history"][-(2 * config.max_history + 1):]

for j, candidate in enumerate(utterance["candidates"][-num_candidates:]):

lm_labels = bool(

j == num_candidates - 1) # the true label is always the last one in list of candidates

instance, _ = build_input_from_segments(persona, history, candidate, tokenizer, lm_labels)

for input_name, input_array in instance.items():

datasets[dataset_name][input_name].append(input_array)

datasets[dataset_name]["mc_labels"].append(num_candidates - 1)

datasets[dataset_name]["n_candidates"] = num_candidates

persona = [persona[-1]] + persona[:-1] # permuted personalities

logger.info("Pad inputs and convert to Tensor")

tensor_datasets = {"train": [], "valid": []}

for dataset_name, dataset in datasets.items():

dataset = pad_dataset(dataset, padding=tokenizer.convert_tokens_to_ids([SPECIAL_TOKENS[-1]])[0])

for input_name in MODEL_INPUTS:

tensor = torch.tensor(dataset[input_name])

if input_name != "mc_labels":

tensor = tensor.view((-1, datasets[dataset_name]["n_candidates"]) + tensor.shape[1:])

tensor_datasets[dataset_name].append(tensor)

logger.info("Build train and validation dataloaders")

train_dataset, valid_dataset = TensorDataset(*tensor_datasets["train"]), TensorDataset(*tensor_datasets["valid"])

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) if config.distributed else None

valid_sampler = torch.utils.data.distributed.DistributedSampler(valid_dataset) if config.distributed else None

train_loader = DataLoader(train_dataset, sampler=train_sampler, batch_size=config.train_batch_size,

shuffle=(not config.distributed))

valid_loader = DataLoader(valid_dataset, sampler=valid_sampler, batch_size=config.valid_batch_size, shuffle=False)

logger.info("Train dataset (Batch, Candidates, Seq length): {}".format(train_dataset.tensors[0].shape))

logger.info("Valid dataset (Batch, Candidates, Seq length): {}".format(valid_dataset.tensors[0].shape))

# parser = ArgumentParser()

# parser.add_argument("--dataset_path", type=str, default="",

# help="Path or url of the dataset. If empty download from S3.")

# parser.add_argument("--dataset_cache", type=str, default='./dataset_cache', help="Path or url of the dataset cache")

# parser.add_argument("--model_checkpoint", type=str,

# default="/home/rohola/codes/transfer-learning-conv-ai/model/bert_large_uncased",

# help="Path, url or short name of the model")

# # parser.add_argument("--model_checkpoint", type=str, default="/home/rohola/codes/transfer-learning-conv-ai/logs/logs4", help="Path, url or short name of the model")

# # parser.add_argument("--model_checkpoint", type=str, default="bert-base-uncased", help="Path, url or short name of the model")

# parser.add_argument("--num_candidates", type=int, default=2, help="Number of candidates for training")

# parser.add_argument("--do_lower_case", default='True', action='store_true')

# parser.add_argument("--max_history", type=int, default=2, help="Number of previous exchanges to keep in history")

# parser.add_argument("--train_batch_size", type=int, default=3, help="Batch size for training")

# parser.add_argument("--valid_batch_size", type=int, default=1, help="Batch size for validation")

# parser.add_argument("--gradient_accumulation_steps", type=int, default=8,

# help="Accumulate gradients on several steps")

# parser.add_argument("--lr", type=float, default=5e-5, help="Learning rate")

# parser.add_argument("--warmup_proportion", type=float, default=0.1, help="Warmup Proportion")

# parser.add_argument("--lm_coef", type=float, default=1.0, help="LM loss coefficient")

# parser.add_argument("--mc_coef", type=float, default=1.0, help="Multiple-choice loss coefficient")

# parser.add_argument("--max_norm", type=float, default=1.0, help="Clipping gradient norm")

# parser.add_argument("--n_epochs", type=int, default=4, help="Number of training epochs")

# parser.add_argument("--personality_permutations", type=int, default=1,

# help="Number of permutations of personality sentences")

# parser.add_argument("--eval_before_start", action='store_true',

# help="If true start with a first evaluation before training")

# parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu",

# help="Device (cuda or cpu)")

# parser.add_argument("--fp16", type=str, default="",

# help="Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)")

# parser.add_argument("--local_rank", type=int, default=-1,

# help="Local rank for distributed training (-1: not distributed)")

# parser.add_argument("--log_dir", type=str, default="",

# help="Local rank for distributed training (-1: not distributed)")

# config = parser.parse_args()

config_file = "configs/train_using_bert_config.json"

config = Config.from_json_file(config_file)

# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes

logging.basicConfig(level=logging.INFO if config.local_rank in [-1, 0] else logging.WARN)

logger.warning("Running process %d",

config.local_rank) # This is a logger.warning: it will be printed by all distributed processes

logger.info("Arguments: %s", pformat(config))

# Initialize distributed training if needed

config.distributed = (config.local_rank != -1)

if config.distributed:

torch.cuda.set_device(config.local_rank)

config.device = torch.device("cuda", config.local_rank)

torch.distributed.init_process_group(backend='nccl', init_method='env://')

logger.info("Prepare tokenizer, pretrained model and optimizer - add special tokens for fine-tuning")

tokenizer = BertTokenizer.from_pretrained(config.model_checkpoint, do_lower_case=config.do_lower_case)

model = BertDoubleHeadsModel.from_pretrained(config.model_checkpoint)

model.to(config.device)

optimizer = BertAdam(model.parameters(), lr=config.lr, warmup=config.warmup_proportion)

# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)

if config.fp16:

from apex import amp # Apex is only required if we use fp16 training

model, optimizer = amp.initialize(model, optimizer, opt_level=config.fp16)

if config.distributed:

model = DistributedDataParallel(model, device_ids=[config.local_rank], output_device=config.local_rank)

logger.info("Prepare datasets")

train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(config, tokenizer)

# Training function and trainer

def update(engine, batch):

model.train()

batch = tuple(input_tensor.to(config.device) for input_tensor in batch)

input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids, input_mask = batch

lm_loss, mc_loss = model(input_ids, mc_token_ids, input_mask, lm_labels, mc_labels, token_type_ids)

loss = (lm_loss * config.lm_coef + mc_loss * config.mc_coef) / config.gradient_accumulation_steps

# mc_loss = model(input_ids, mc_token_ids, input_mask, lm_labels=None, mc_labels=mc_labels, token_type_ids=token_type_ids)

# loss = mc_loss[0] / config.gradient_accumulation_steps

if config.fp16:

with amp.scale_loss(loss, optimizer) as scaled_loss:

scaled_loss.backward()

torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), config.max_norm)

else:

loss.backward()

torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_norm)

if engine.state.iteration % config.gradient_accumulation_steps == 0:

optimizer.step()

optimizer.zero_grad()

return loss.item()

trainer = Engine(update)

# Evaluation function and evaluator (evaluator output is the input of the metrics)

def inference(engine, batch):

model.eval()

with torch.no_grad():

batch = tuple(input_tensor.to(config.device) for input_tensor in batch)

input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids, input_mask = batch

model_outputs = model(input_ids, mc_token_ids, input_mask, token_type_ids=token_type_ids)

lm_logits, mc_logits = model_outputs[0], model_outputs[1]

lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(-1, lm_logits.size(-1))

lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)

# loss_fct = torch.nn.CrossEntropyLoss(ignore_index=-1)

# l = loss_fct(lm_logits_flat_shifted, lm_labels_flat_shifted)

return (lm_logits_flat_shifted, mc_logits), (lm_labels_flat_shifted, mc_labels)

evaluator = Engine(inference)

# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch

trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda _: evaluator.run(val_loader))

if config.n_epochs < 1:

trainer.add_event_handler(Events.COMPLETED, lambda _: evaluator.run(val_loader))

if config.eval_before_start:

trainer.add_event_handler(Events.STARTED, lambda _: evaluator.run(val_loader))

# Make sure distributed data samplers split the dataset nicely between the distributed processes

if config.distributed:

trainer.add_event_handler(Events.EPOCH_STARTED, lambda engine: train_sampler.set_epoch(engine.state.epoch))

evaluator.add_event_handler(Events.EPOCH_STARTED, lambda engine: valid_sampler.set_epoch(engine.state.epoch))

# Linearly decrease the learning rate from lr to zero

scheduler = PiecewiseLinear(optimizer, "lr", [(0, config.lr), (config.n_epochs * len(train_loader), 0.0)])

trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)

# Prepare metrics - note how we compute distributed metrics

RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")

metrics = {"nll": Loss(torch.nn.CrossEntropyLoss(ignore_index=-1), output_transform=lambda x: (x[0][0], x[1][0])),

"accuracy": Accuracy(output_transform=lambda x: (x[0][1], x[1][1]))}

metrics.update({"average_nll": MetricsLambda(average_distributed_scalar, metrics["nll"], config),

"average_accuracy": MetricsLambda(average_distributed_scalar, metrics["accuracy"], config)})

metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])

for name, metric in metrics.items():

metric.attach(evaluator, name)

# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train

if config.local_rank in [-1, 0]:

pbar = ProgressBar(persist=True)

pbar.attach(trainer, metric_names=["loss"])

evaluator.add_event_handler(Events.COMPLETED,

lambda _: pbar.log_message("Validation: %s" % pformat(evaluator.state.metrics)))

tb_logger = TensorboardLogger(log_dir=config.log_dir)

tb_logger.attach(trainer, log_handler=OutputHandler(tag="training", metric_names=["loss"]),

event_name=Events.ITERATION_COMPLETED)

tb_logger.attach(trainer, log_handler=OptimizerParamsHandler(optimizer), event_name=Events.ITERATION_STARTED)

tb_logger.attach(evaluator, log_handler=OutputHandler(tag="validation", metric_names=list(metrics.keys()),

another_engine=trainer),

event_name=Events.EPOCH_COMPLETED)

checkpoint_handler = ModelCheckpoint(tb_logger.writer.log_dir, 'checkpoint', save_interval=1, n_saved=3)

trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler, {

'mymodel': getattr(model, 'module', model)}) # "getattr" take care of distributed encapsulation

torch.save(config, tb_logger.writer.log_dir + '/model_training_args.bin')

getattr(model, 'module', model).config.to_json_file(os.path.join(tb_logger.writer.log_dir, CONFIG_NAME))

tokenizer.save_vocabulary(tb_logger.writer.log_dir)

# Run the training

trainer.run(train_loader, max_epochs=config.n_epochs)

# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)

if config.local_rank in [-1, 0] and config.n_epochs > 0:

os.rename(checkpoint_handler._saved[-1][1][-1], os.path.join(tb_logger.writer.log_dir,

WEIGHTS_NAME)) # TODO: PR in ignite to have better access to saved file paths (cleaner)