text: from datetime import datetime
  
  import numpy as np
  import random
  import os
  import sys
  import json
  import shutil
  import torch
  import torch.nn as nn
  import torch.distributed as dist
  from torch.utils.data import DataLoader, Dataset
  from torch.utils.data.sampler import RandomSampler
  from torch.utils.data.distributed import DistributedSampler
  
  import argparse
  from tqdm import tqdm
  from checkpoint import checkpoint_model, load_checkpoint, latest_checkpoint_file
  from logger import Logger
  from utils import get_sample_writer
  from models import BertMultiTask
  from dataset import PreTrainingDataset
  from dataset import PretrainDataType
  from pytorch_pretrained_bert.tokenization import BertTokenizer
  from pytorch_pretrained_bert.optimization import BertAdam
  from optimization import warmup_linear_decay_exp
  from azureml_adapter import set_environment_variables_for_nccl_backend, get_local_rank, get_global_size, get_local_size
  from sources import PretrainingDataCreator, TokenInstance, GenericPretrainingDataCreator
  from sources import WikiPretrainingDataCreator
  from configuration import BertJobConfiguration
  
  from azureml.core.run import Run
  
  
  def get_effective_batch(total):
      if use_multigpu_with_single_device_per_process:
          return total//dist.get_world_size()//train_batch_size//gradient_accumulation_steps
      else:
          return total//train_batch_size//gradient_accumulation_steps # Dividing with gradient_accumulation_steps since we multiplied it earlier
  
  
  def get_dataloader(dataset: Dataset, eval_set=False):
      if not use_multigpu_with_single_device_per_process:
          train_sampler = RandomSampler(dataset)
      else:
          train_sampler = DistributedSampler(dataset)
      return (x for x in DataLoader(dataset, batch_size=train_batch_size // 2 if eval_set else train_batch_size,
                                    sampler=train_sampler, num_workers=job_config.get_num_workers()))
  
  
  def pretrain_validation(index):
      model.eval()
      dataset = PreTrainingDataset(tokenizer=tokenizer,
                                   folder=args.validation_path,
                                   logger=logger, max_seq_length=max_seq_length,
                                   index=index, data_type=PretrainDataType.VALIDATION,
                                   max_predictions_per_seq=max_predictions_per_seq,
                                   masked_lm_prob=masked_lm_prob)
      data_batches = get_dataloader(dataset, eval_set=True)
      eval_loss = 0
      nb_eval_steps = 0
  
      for batch in data_batches:
          batch = tuple(t.to(device) for t in batch)
          tmp_eval_loss = model.network(batch, log=False)
          dist.reduce(tmp_eval_loss, 0)
          # Reduce to get the loss from all the GPU's
          tmp_eval_loss = tmp_eval_loss / dist.get_world_size()
          eval_loss += tmp_eval_loss.mean().item()
          nb_eval_steps += 1
      eval_loss = eval_loss / nb_eval_steps
      logger.info(f"Validation Loss for epoch {index + 1} is: {eval_loss}")
      if check_write_log():
          summary_writer.add_scalar(f'Validation/Loss', eval_loss, index + 1)
          run.log("validation_loss", np.float(eval_loss))
          run.log_row("validation_loss over epochs", epoch = index, val_loss = np.float(eval_loss))
      return eval_loss
  
  
  def train(index):
      model.train()
      dataloaders = {}
      i = 0
      global global_step
      datalengths = []
      batchs_per_dataset = []
  
      # Pretraining datasets
      wiki_pretrain_dataset = PreTrainingDataset(tokenizer=tokenizer,
                                                 folder=args.train_path,
                                                 logger=logger, max_seq_length=max_seq_length,
                                                 index=index, data_type=PretrainDataType.WIKIPEDIA,
                                                 max_predictions_per_seq=max_predictions_per_seq,
                                                 masked_lm_prob=masked_lm_prob)
  
      datalengths.append(len(wiki_pretrain_dataset))
      dataloaders[i] = get_dataloader(wiki_pretrain_dataset)
  
      num_batches_in_dataset = get_effective_batch(len(wiki_pretrain_dataset))
      logger.info('Wikpedia data file: Number of samples {}, number of batches required to process these samples: {}'.format(len(wiki_pretrain_dataset), num_batches_in_dataset))
      
      batchs_per_dataset.append(num_batches_in_dataset)
      i += 1
  
      logger.info("Training on Wikipedia dataset")
  
      total_length = sum(datalengths)
  
      dataset_batches = []
      for i, batch_count in enumerate(batchs_per_dataset):
          dataset_batches.extend([i] * batch_count)
      logger.info('Number of batches to process *all* data samples in this epoch: {}'.format(len(dataset_batches)))
      # shuffle
      random.shuffle(dataset_batches)
  
      # We don't want the dataset to be n the form of alternate chunks if we have more than
      # one dataset type, instead we want to organize them into contiguous chunks of each
      # data type, hence the multiplication with grad_accumulation_steps with dataset_batch_type
      dataset_picker = []
      for dataset_batch_type in dataset_batches:
          dataset_picker.extend([dataset_batch_type] * gradient_accumulation_steps )
  
      logger.info('Number of steps to process all batches in this epoch: {}'.format(len(dataset_picker)))
      model.train()
  
      # Counter of sequences in an "epoch"
      sequences_counter = 0
      global_step_loss = 0
  
      for step, dataset_type in enumerate(dataset_picker):
          try:
              batch = next(dataloaders[dataset_type])
  
              sequences_counter += len(batch)
  
              if n_gpu == 1:
                  batch = tuple(t.to(device) for t in batch)  # Move to GPU
  
              if step > 1 and step % 1000 == 0:
                  logger.info("{} Number of sequences processed so far: {} (cumulative in {} steps)".format(datetime.utcnow(), sequences_counter, step))
              # Calculate forward pass
              loss = model.network(batch)
  
              if n_gpu > 1:
                  # this is to average loss for multi-gpu. In DistributedDataParallel
                  # setting, we get tuple of losses form all proccesses
                  loss = loss.mean()
  
              if gradient_accumulation_steps > 1:
                  loss = loss / gradient_accumulation_steps
  
              # Enabling  optimized Reduction
              # reduction only happens in backward if this method is called before
              # when using the distributed module
              if accumulate_gradients:
                  if use_multigpu_with_single_device_per_process and (step + 1) % gradient_accumulation_steps == 0:
                      model.network.enable_need_reduction()
                  else:
                      model.network.disable_need_reduction()
              if fp16:
                  optimizer.backward(loss)
              else:
                  loss.backward()
  
              global_step_loss += loss
              if (step + 1) % gradient_accumulation_steps == 0:
                  if fp16:
                      # modify learning rate with special warm up BERT uses
                      # if fp16 is False, BertAdam is used that handles this automatically
                      lr_this_step = \
                          job_config.get_learning_rate() * warmup_linear_decay_exp(global_step,
                                                                                   job_config.get_decay_rate(),
                                                                                   job_config.get_decay_step(),
                                                                                   job_config.get_total_training_steps(),
                                                                                   job_config.get_warmup_proportion())
                      for param_group in optimizer.param_groups:
                          param_group['lr'] = lr_this_step
  
                      # Record the LR against global_step on tensorboard
                      if check_write_log():
                          summary_writer.add_scalar(f'Train/lr', lr_this_step, global_step)
                      
                  optimizer.step()
                  optimizer.zero_grad()
                  global_step += 1
                  if check_write_log() and (global_step%args.log_steps == 0):
                      run.log("training_loss", np.float(global_step_loss))
                      run.log("lr_this_step", np.float(lr_this_step))
                      run.log_row("loss over steps", global_step = global_step, loss =  np.float(global_step_loss))
                      run.log_row("lr over steps", global_step = global_step, lr  = np.float(lr_this_step))
                  global_step_loss = 0
          except StopIteration:
              continue
          
      logger.info("Completed {} steps".format(step))
      logger.info("Completed processing {} sequences".format(sequences_counter))
  
      # Run Validation Loss
      if max_seq_length == 512:
          logger.info(f"TRAIN BATCH SIZE: {train_batch_size}")
          return pretrain_validation(index)
      else:
          return None
  
  
  def str2bool(val):
      return val.lower() == "true" or val.lower() == "t" or val.lower() == "1"
  
  def check_write_log():
      return dist.get_rank() == 0 or not use_multigpu_with_single_device_per_process
  
  if __name__ == '__main__':
      print("The arguments are: " + str(sys.argv))
  
      parser = argparse.ArgumentParser()
  
      # Required_parameters
      parser.add_argument("--config_file", "--cf",
                          help="pointer to the configuration file of the experiment", type=str, required=True)
  
      parser.add_argument("--config_file_path", default=None, type=str, required=True,
                          help="The blob storage directory where config file is located.")
  
      parser.add_argument("--train_path", default=None, type=str, required=True,
                          help="The blob storage directory for train data, cache and output.")
  
      parser.add_argument("--validation_path", default=None, type=str, required=True,
                          help="The blob storage directory for validation data, cache and output.")
  
      parser.add_argument('--tokenizer_path', type=str, default=False,
                      help="Path to load the tokenizer from")
      parser.add_argument("--output_dir", default=None, type=str, required=True,
                          help="If given, model checkpoints will be saved to this directory.")
      
      # Optional Params
      parser.add_argument("--best_cp_dir", default=None, type=str,
                          help="If given, model best checkpoint will be saved to this directory.")
      parser.add_argument("--latest_cp_dir", default=None, type=str,
                          help="If given, model latest checkpoint will be saved to this directory.")
      parser.add_argument("--max_seq_length", default=512, type=int,
                          help="The maximum total input sequence length after WordPiece tokenization. Sequences "
                               "longer than this will be truncated, and sequences shorter than this will be padded.")
      parser.add_argument("--max_predictions_per_seq", "--max_pred", default=80, type=int,
                          help="The maximum number of masked tokens in a sequence to be predicted.")
      parser.add_argument("--masked_lm_prob", "--mlm_prob", default=0.15,
                          type=float, help="The masking probability for languge model.")
      parser.add_argument("--train_batch_size", default=32,
                          type=int, help="Total batch size for training.")
      parser.add_argument("--no_cuda",
                          type=str,
                          default='False',
                          help="Whether not to use CUDA when available")
      parser.add_argument('--seed',
                          type=int,
                          default=42,
                          help="random seed for initialization")
      parser.add_argument('--accumulate_gradients',
                          type=str,
                          default='True',
                          help="Enabling gradient accumulation optimization")
      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('--fp16',
                          type=str,
                          default='False',
                          help="Whether to use 16-bit float precision instead of 32-bit")
      parser.add_argument('--use_pretrain',
                          type=str,
                          default='False',
                          help="Whether to use Bert Pretrain Weights or not")
      parser.add_argument('--loss_scale',
                          type=float,
                          default=0,
                          help='Loss scaling, positive power of 2 values can improve fp16 convergence.')
      parser.add_argument('--load_training_checkpoint', '--load_cp',
                          type=str,
                          default='False',
                          help="This is the path to the TAR file which contains model+opt state_dict() checkpointed.")
      parser.add_argument('--use_multigpu_with_single_device_per_process',
                          type=str,
                          default='True',
                          help="Whether only one device is managed per process")	    
      parser.add_argument('--epochs',		
                          type=int,		
                          default=250,		
                          help="total number of epochs")
      parser.add_argument('--log_steps',		
                          type=int,		
                          default=50,		
                          help="logging intervals")
      parser.add_argument('--backend',		
                          type=str,		
                          default='nccl',		
                          help="reduce backend to use")
      
      args = parser.parse_args()
  
      if args.output_dir:
          os.makedirs(args.output_dir, exist_ok=True)
      if args.best_cp_dir:
          os.makedirs(args.best_cp_dir, exist_ok=True)
      if args.latest_cp_dir:
          os.makedirs(args.latest_cp_dir, exist_ok=True)
  
      no_cuda = str2bool(args.no_cuda)
      fp16 = str2bool(args.fp16)
      accumulate_gradients = str2bool(args.accumulate_gradients)
      use_pretrain = str2bool(args.use_pretrain)
      use_multigpu_with_single_device_per_process = str2bool(args.use_multigpu_with_single_device_per_process)
  
      config_file = args.config_file
      gradient_accumulation_steps = args.gradient_accumulation_steps
      train_batch_size = args.train_batch_size
      seed = args.seed
      loss_scale = args.loss_scale
      load_training_checkpoint = args.load_training_checkpoint
      max_seq_length = args.max_seq_length
      max_predictions_per_seq = args.max_predictions_per_seq
      masked_lm_prob = args.masked_lm_prob
  
      local_rank = -1
  
      local_rank = get_local_rank()
      global_size = get_global_size()
      local_size = get_local_size()	
      # TODO use logger	
      print('local_rank = {}'.format(local_rank))
      print('global_size = {}'.format(global_size))
      print('local_size = {}'.format(local_size))
  
      set_environment_variables_for_nccl_backend(local_size == global_size)
  
      # Prepare Logger
      logger = Logger(cuda=torch.cuda.is_available())
  
      # # Extact config file from blob storage
      job_config = BertJobConfiguration(config_file_path=os.path.join(args.config_file_path, config_file))
  
      job_name = job_config.get_name()
      # Setting the distributed variables
  
      run = Run.get_context()
  
      if not use_multigpu_with_single_device_per_process:
          device = torch.device("cuda")
          n_gpu = torch.cuda.device_count()
      else:
          device = torch.device("cuda", local_rank)
          n_gpu = 1
          # Initializes the distributed backend which will take care of synchronizing nodes/GPUs
          torch.distributed.init_process_group(backend=args.backend)
          if fp16:
              logger.info("16-bits distributed training is not officially supported in the version of PyTorch currently used, but it works. Refer to https://github.com/pytorch/pytorch/pull/13496 for supported version.")
              fp16 = True  #
      logger.info("device: {} n_gpu: {}, use_multigpu_with_single_device_per_process: {}, 16-bits training: {}".format(
          device, n_gpu, use_multigpu_with_single_device_per_process, fp16))
  
      if gradient_accumulation_steps < 1:
          raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
              gradient_accumulation_steps))
  
      train_batch_size = int(train_batch_size / gradient_accumulation_steps)
  
      # Setting all the seeds so that the task is random but same accross processes
      random.seed(seed)
      np.random.seed(seed)
      torch.manual_seed(seed)
      logger.info
      if n_gpu > 0:
          torch.cuda.manual_seed_all(seed)
  
      # Create an outputs/ folder in the blob storage
      if args.output_dir is None:
          parent_dir = os.path.join(args.output_dir, 'outputs', str(run.experiment.name))
          output_dir = os.path.join(parent_dir, str(run.id))
          os.makedirs(output_dir, exist_ok=True)
          saved_model_path = os.path.join(output_dir, "saved_models", job_name)
          os.makedirs(saved_model_path, exist_ok=True)
      else:
          saved_model_path = args.output_dir
  
      summary_writer = None
      # Prepare Summary Writer and saved_models path
      if check_write_log():
          #azureml.tensorboard only streams from /logs directory, therefore hardcoded
          summary_writer = get_sample_writer(
              name=job_name, base='./logs')
  
      # Loading Tokenizer (vocabulary from blob storage, if exists)
      logger.info("Extracting the vocabulary")
      if args.tokenizer_path:
          logger.info(f'Loading tokenizer from {args.tokenizer_path}')
          tokenizer = BertTokenizer.from_pretrained(
              args.tokenizer_path, cache_dir=args.output_dir)
      else:
          tokenizer = BertTokenizer.from_pretrained(job_config.get_token_file_type(), cache_dir=args.output_dir)
      logger.info("Vocabulary contains {} tokens".format(len(list(tokenizer.vocab.keys()))))
  
  
      # Loading Model
      logger.info("Initializing BertMultiTask model")
      model = BertMultiTask(job_config = job_config, use_pretrain = use_pretrain, tokenizer = tokenizer, 
                            cache_dir = args.output_dir, device = device, write_log = check_write_log(), 
                            summary_writer = summary_writer)
  
      logger.info("Converting the input parameters")
      if fp16:
          model.half()
          
      model.to(device)
  
      if use_multigpu_with_single_device_per_process:
          try:
              if accumulate_gradients:
                  logger.info("Enabling gradient accumulation by using a forked version of DistributedDataParallel implementation available in the branch bertonazureml/apex at https://www.github.com/microsoft/apex")
                  from distributed_apex import DistributedDataParallel as DDP
              else:
                  logger.info("Using Default Apex DistributedDataParallel implementation")
                  from apex.parallel import DistributedDataParallel as DDP
          except ImportError:
              raise ImportError("To use distributed and fp16 training, please install apex from the branch bertonazureml/apex at https://www.github.com/microsoft/apex.")
          torch.cuda.set_device(local_rank)
          model.network = DDP(model.network, delay_allreduce=False)
  
      elif n_gpu > 1:
          model.network = nn.DataParallel(model.network)
  
      # Prepare Optimizer
      logger.info("Preparing the optimizer")
      param_optimizer = list(model.network.named_parameters())
      param_optimizer = [n for n in param_optimizer if 'pooler' not in n[0]]
      no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
      optimizer_grouped_parameters = [
          {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
          {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
      ]
  
      logger.info("Loading Apex and building the FusedAdam optimizer")
  
      if fp16:
          try:
              from apex.optimizers import FP16_Optimizer, FusedAdam
          except:
              raise ImportError("To use distributed and fp16 training, please install apex from the branch bertonazureml/apex at https://www.github.com/microsoft/apex.")
  
          optimizer = FusedAdam(optimizer_grouped_parameters,
                                lr=job_config.get_learning_rate(),
                                bias_correction=False,
                                max_grad_norm=1.0)
          if loss_scale == 0:
              optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
          else:
              optimizer = FP16_Optimizer(
                  optimizer, static_loss_scale=loss_scale)
      else:
          optimizer = BertAdam(optimizer_grouped_parameters,
                               lr=job_config.get_learning_rate(),
                               warmup=job_config.get_warmup_proportion(),
                               t_total=job_config.get_total_training_steps())
  
      global_step = 0
      start_epoch = 0
      
      # if args.load_training_checkpoint is not None:
      if load_training_checkpoint != 'False':
          logger.info(f"Looking for previous training checkpoint.")
          latest_checkpoint_path = latest_checkpoint_file(args.load_training_checkpoint, no_cuda)
  
          logger.info(f"Restoring previous training checkpoint from {latest_checkpoint_path}")
          start_epoch, global_step = load_checkpoint(model, optimizer, latest_checkpoint_path)
          logger.info(f"The model is loaded from last checkpoint at epoch {start_epoch} when the global steps were at {global_step}")
  
  
      logger.info("Training the model")
  
      best_loss = None
      for index in range(start_epoch, args.epochs):
          logger.info(f"Training epoch: {index + 1}")
          
          eval_loss = train(index)
  
          if check_write_log():
              if best_loss is None or eval_loss is None or eval_loss < best_loss*0.99:
                  best_loss = eval_loss
                  epoch_ckp_path = os.path.join(saved_model_path, "bert_encoder_epoch_{0:04d}.pt".format(index + 1))
                  checkpoint_model(os.path.join(saved_model_path, "training_state_checkpoint_{0:04d}.tar".format(index + 1)), model, optimizer, index, global_step)
                  logger.info(f"Saving checkpoint of the model from epoch {index + 1} at {epoch_ckp_path}")
                  model.save_bert(epoch_ckp_path)
  
                  #save best checkpoint in separate directory
                  if args.best_cp_dir:
                      best_ckp_path = os.path.join(args.best_cp_dir, "bert_encoder_epoch_{0:04d}.pt".format(index + 1))
                      shutil.rmtree(args.best_cp_dir)
                      os.makedirs(args.best_cp_dir,exist_ok=True)
                      model.save_bert(best_ckp_path)
                  
              if args.latest_cp_dir:
                  shutil.rmtree(args.latest_cp_dir)
                  os.makedirs(args.latest_cp_dir,exist_ok=True)
                  checkpoint_model(os.path.join(args.latest_cp_dir, "training_state_checkpoint_{0:04d}.tar".format(index + 1)), model, optimizer, index, global_step)
                  latest_ckp_path = os.path.join(args.latest_cp_dir, "bert_encoder_epoch_{0:04d}.pt".format(index + 1))
                  model.save_bert(latest_ckp_path)
  