code: def main():
      args = parse_args()
  
      # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
      # information sent is the one passed as arguments along with your Python/PyTorch versions.
      send_example_telemetry("run_qa_beam_search_no_trainer", args)
  
      # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
      # If we're using tracking, we also need to initialize it here and it will pick up all supported trackers
      # in the environment
      accelerator_log_kwargs = {}
  
      if args.with_tracking:
          accelerator_log_kwargs["log_with"] = args.report_to
          accelerator_log_kwargs["logging_dir"] = args.output_dir
  
      accelerator = Accelerator(
          gradient_accumulation_steps=args.gradient_accumulation_steps,
          **accelerator_log_kwargs,
      )
  
      # Make one log on every process with the configuration for debugging.
      logging.basicConfig(
          format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
          datefmt="%m/%d/%Y %H:%M:%S",
          level=logging.INFO,
      )
      logger.info(accelerator.state, main_process_only=False)
      if accelerator.is_local_main_process:
          datasets.utils.logging.set_verbosity_warning()
          transformers.utils.logging.set_verbosity_info()
      else:
          datasets.utils.logging.set_verbosity_error()
          transformers.utils.logging.set_verbosity_error()
  
      # If passed along, set the training seed now.
      if args.seed is not None:
          set_seed(args.seed)
  
      # Handle the repository creation
      if accelerator.is_main_process:
          if args.push_to_hub:
              if args.hub_model_id is None:
                  repo_name = get_full_repo_name(
                      Path(args.output_dir).name, token=args.hub_token
                  )
              else:
                  repo_name = args.hub_model_id
              repo = Repository(args.output_dir, clone_from=repo_name)
  
              with open(
                  os.path.join(args.output_dir, ".gitignore"), "w+"
              ) as gitignore:
                  if "step_*" not in gitignore:
                      gitignore.write("step_*\n")
                  if "epoch_*" not in gitignore:
                      gitignore.write("epoch_*\n")
          elif args.output_dir is not None:
              os.makedirs(args.output_dir, exist_ok=True)
      accelerator.wait_for_everyone()
  
      # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
      # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
      # (the dataset will be downloaded automatically from the datasets Hub).
      #
      # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
      # 'text' is found. You can easily tweak this behavior (see below).
      #
      # In distributed training, the load_dataset function guarantee that only one local process can concurrently
      # download the dataset.
      if args.dataset_name is not None:
          # Downloading and loading a dataset from the hub.
          raw_datasets = load_dataset(
              args.dataset_name, args.dataset_config_name
          )
      else:
          data_files = {}
          if args.train_file is not None:
              data_files["train"] = args.train_file
          if args.validation_file is not None:
              data_files["validation"] = args.validation_file
          if args.test_file is not None:
              data_files["test"] = args.test_file
          extension = args.train_file.split(".")[-1]
          raw_datasets = load_dataset(
              extension, data_files=data_files, field="data"
          )
      # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
      # https://huggingface.co/docs/datasets/loading_datasets.html.
  
      # Load pretrained model and tokenizer
      #
      # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently
      # download model & vocab.
  
      config = XLNetConfig.from_pretrained(args.model_name_or_path)
      tokenizer = XLNetTokenizerFast.from_pretrained(
          args.model_name_or_path
      )
      model = XLNetForQuestionAnswering.from_pretrained(
          args.model_name_or_path,
          from_tf=bool(".ckpt" in args.model_name_or_path),
          config=config,
      )
  
      # Preprocessing the datasets.
      # Preprocessing is slighlty different for training and evaluation.
      column_names = raw_datasets["train"].column_names
  
      question_column_name = (
          "question" if "question" in column_names else column_names[0]
      )
      context_column_name = (
          "context" if "context" in column_names else column_names[1]
      )
      answer_column_name = (
          "answers" if "answers" in column_names else column_names[2]
      )
  
      # Padding side determines if we do (question|context) or (context|question).
      pad_on_right = tokenizer.padding_side == "right"
  
      if args.max_seq_length > tokenizer.model_max_length:
          logger.warning(
              f"The max_seq_length passed ({args.max_seq_length}) is larger than the maximum length for the"
              f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
          )
  
      max_seq_length = min(
          args.max_seq_length, tokenizer.model_max_length
      )
  
      # Training preprocessing
      def prepare_train_features(examples):
          # Some of the questions have lots of whitespace on the left, which is not useful and will make the
          # truncation of the context fail (the tokenized question will take a lots of space). So we remove that
          # left whitespace
          examples[question_column_name] = [
              q.lstrip() for q in examples[question_column_name]
          ]
  
          # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results
          # in one example possible giving several features when a context is long, each of those features having a
          # context that overlaps a bit the context of the previous feature.
          tokenized_examples = tokenizer(
              examples[
                  question_column_name
                  if pad_on_right
                  else context_column_name
              ],
              examples[
                  context_column_name
                  if pad_on_right
                  else question_column_name
              ],
              truncation="only_second"
              if pad_on_right
              else "only_first",
              max_length=max_seq_length,
              stride=args.doc_stride,
              return_overflowing_tokens=True,
              return_offsets_mapping=True,
              return_special_tokens_mask=True,
              return_token_type_ids=True,
              padding="max_length",
          )
  
          # Since one example might give us several features if it has a long context, we need a map from a feature to
          # its corresponding example. This key gives us just that.
          sample_mapping = tokenized_examples.pop(
              "overflow_to_sample_mapping"
          )
          # The offset mappings will give us a map from token to character position in the original context. This will
          # help us compute the start_positions and end_positions.
          offset_mapping = tokenized_examples.pop("offset_mapping")
          # The special tokens will help us build the p_mask (which indicates the tokens that can't be in answers).
          special_tokens = tokenized_examples.pop("special_tokens_mask")
  
          # Let's label those examples!
          tokenized_examples["start_positions"] = []
          tokenized_examples["end_positions"] = []
          tokenized_examples["is_impossible"] = []
          tokenized_examples["cls_index"] = []
          tokenized_examples["p_mask"] = []
  
          for i, offsets in enumerate(offset_mapping):
              # We will label impossible answers with the index of the CLS token.
              input_ids = tokenized_examples["input_ids"][i]
              cls_index = input_ids.index(tokenizer.cls_token_id)
              tokenized_examples["cls_index"].append(cls_index)
  
              # Grab the sequence corresponding to that example (to know what is the context and what is the question).
              sequence_ids = tokenized_examples["token_type_ids"][i]
              for k, s in enumerate(special_tokens[i]):
                  if s:
                      sequence_ids[k] = 3
              context_idx = 1 if pad_on_right else 0
  
              # Build the p_mask: non special tokens and context gets 0.0, the others get 1.0.
              # The cls token gets 1.0 too (for predictions of empty answers).
              tokenized_examples["p_mask"].append(
                  [
                      0.0
                      if (not special_tokens[i][k] and s == context_idx)
                      or k == cls_index
                      else 1.0
                      for k, s in enumerate(sequence_ids)
                  ]
              )
  
              # One example can give several spans, this is the index of the example containing this span of text.
              sample_index = sample_mapping[i]
              answers = examples[answer_column_name][sample_index]
              # If no answers are given, set the cls_index as answer.
              if len(answers["answer_start"]) == 0:
                  tokenized_examples["start_positions"].append(
                      cls_index
                  )
                  tokenized_examples["end_positions"].append(cls_index)
                  tokenized_examples["is_impossible"].append(1.0)
              else:
                  # Start/end character index of the answer in the text.
                  start_char = answers["answer_start"][0]
                  end_char = start_char + len(answers["text"][0])
  
                  # Start token index of the current span in the text.
                  token_start_index = 0
                  while sequence_ids[token_start_index] != context_idx:
                      token_start_index += 1
  
                  # End token index of the current span in the text.
                  token_end_index = len(input_ids) - 1
                  while sequence_ids[token_end_index] != context_idx:
                      token_end_index -= 1
                  # Detect if the answer is out of the span (in which case this feature is labeled with the CLS index).
                  if not (
                      offsets[token_start_index][0] <= start_char
                      and offsets[token_end_index][1] >= end_char
                  ):
                      tokenized_examples["start_positions"].append(
                          cls_index
                      )
                      tokenized_examples["end_positions"].append(
                          cls_index
                      )
                      tokenized_examples["is_impossible"].append(1.0)
                  else:
                      # Otherwise move the token_start_index and token_end_index to the two ends of the answer.
                      # Note: we could go after the last offset if the answer is the last word (edge case).
                      while (
                          token_start_index < len(offsets)
                          and offsets[token_start_index][0]
                          <= start_char
                      ):
                          token_start_index += 1
                      tokenized_examples["start_positions"].append(
                          token_start_index - 1
                      )
                      while offsets[token_end_index][1] >= end_char:
                          token_end_index -= 1
                      tokenized_examples["end_positions"].append(
                          token_end_index + 1
                      )
                      tokenized_examples["is_impossible"].append(0.0)
  
          return tokenized_examples
  
      if "train" not in raw_datasets:
          raise ValueError("--do_train requires a train dataset")
      train_dataset = raw_datasets["train"]
      if args.max_train_samples is not None:
          # We will select sample from whole data if agument is specified
          train_dataset = train_dataset.select(
              range(args.max_train_samples)
          )
      # Create train feature from dataset
      with accelerator.main_process_first():
          train_dataset = train_dataset.map(
              prepare_train_features,
              batched=True,
              num_proc=args.preprocessing_num_workers,
              remove_columns=column_names,
              load_from_cache_file=not args.overwrite_cache,
              desc="Running tokenizer on train dataset",
          )
      if args.max_train_samples is not None:
          # Number of samples might increase during Feature Creation, We select only specified max samples
          train_dataset = train_dataset.select(
              range(args.max_train_samples)
          )
  
      # Validation preprocessing
      def prepare_validation_features(examples):
          # Some of the questions have lots of whitespace on the left, which is not useful and will make the
          # truncation of the context fail (the tokenized question will take a lots of space). So we remove that
          # left whitespace
          examples[question_column_name] = [
              q.lstrip() for q in examples[question_column_name]
          ]
  
          # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results
          # in one example possible giving several features when a context is long, each of those features having a
          # context that overlaps a bit the context of the previous feature.
          tokenized_examples = tokenizer(
              examples[
                  question_column_name
                  if pad_on_right
                  else context_column_name
              ],
              examples[
                  context_column_name
                  if pad_on_right
                  else question_column_name
              ],
              truncation="only_second"
              if pad_on_right
              else "only_first",
              max_length=max_seq_length,
              stride=args.doc_stride,
              return_overflowing_tokens=True,
              return_offsets_mapping=True,
              return_special_tokens_mask=True,
              return_token_type_ids=True,
              padding="max_length",
          )
  
          # Since one example might give us several features if it has a long context, we need a map from a feature to
          # its corresponding example. This key gives us just that.
          sample_mapping = tokenized_examples.pop(
              "overflow_to_sample_mapping"
          )
  
          # The special tokens will help us build the p_mask (which indicates the tokens that can't be in answers).
          special_tokens = tokenized_examples.pop("special_tokens_mask")
  
          # For evaluation, we will need to convert our predictions to substrings of the context, so we keep the
          # corresponding example_id and we will store the offset mappings.
          tokenized_examples["example_id"] = []
  
          # We still provide the index of the CLS token and the p_mask to the model, but not the is_impossible label.
          tokenized_examples["cls_index"] = []
          tokenized_examples["p_mask"] = []
  
          for i, input_ids in enumerate(
              tokenized_examples["input_ids"]
          ):
              # Find the CLS token in the input ids.
              cls_index = input_ids.index(tokenizer.cls_token_id)
              tokenized_examples["cls_index"].append(cls_index)
  
              # Grab the sequence corresponding to that example (to know what is the context and what is the question).
              sequence_ids = tokenized_examples["token_type_ids"][i]
              for k, s in enumerate(special_tokens[i]):
                  if s:
                      sequence_ids[k] = 3
              context_idx = 1 if pad_on_right else 0
  
              # Build the p_mask: non special tokens and context gets 0.0, the others 1.0.
              tokenized_examples["p_mask"].append(
                  [
                      0.0
                      if (not special_tokens[i][k] and s == context_idx)
                      or k == cls_index
                      else 1.0
                      for k, s in enumerate(sequence_ids)
                  ]
              )
  
              # One example can give several spans, this is the index of the example containing this span of text.
              sample_index = sample_mapping[i]
              tokenized_examples["example_id"].append(
                  examples["id"][sample_index]
              )
  
              # Set to None the offset_mapping that are not part of the context so it's easy to determine if a token
              # position is part of the context or not.
              tokenized_examples["offset_mapping"][i] = [
                  (o if sequence_ids[k] == context_idx else None)
                  for k, o in enumerate(
                      tokenized_examples["offset_mapping"][i]
                  )
              ]
  
          return tokenized_examples
  
      if "validation" not in raw_datasets:
          raise ValueError("--do_eval requires a validation dataset")
      eval_examples = raw_datasets["validation"]
      if args.max_eval_samples is not None:
          # We will select sample from whole data
          eval_examples = eval_examples.select(
              range(args.max_eval_samples)
          )
      # Validation Feature Creation
      with accelerator.main_process_first():
          eval_dataset = eval_examples.map(
              prepare_validation_features,
              batched=True,
              num_proc=args.preprocessing_num_workers,
              remove_columns=column_names,
              load_from_cache_file=not args.overwrite_cache,
              desc="Running tokenizer on validation dataset",
          )
  
      if args.max_eval_samples is not None:
          # During Feature creation dataset samples might increase, we will select required samples again
          eval_dataset = eval_dataset.select(
              range(args.max_eval_samples)
          )
  
      if args.do_predict:
          if "test" not in raw_datasets:
              raise ValueError("--do_predict requires a test dataset")
          predict_examples = raw_datasets["test"]
          if args.max_predict_samples is not None:
              # We will select sample from whole data
              predict_examples = predict_examples.select(
                  range(args.max_predict_samples)
              )
          # Predict Feature Creation
          with accelerator.main_process_first():
              predict_dataset = predict_examples.map(
                  prepare_validation_features,
                  batched=True,
                  num_proc=args.preprocessing_num_workers,
                  remove_columns=column_names,
                  load_from_cache_file=not args.overwrite_cache,
                  desc="Running tokenizer on prediction dataset",
              )
              if args.max_predict_samples is not None:
                  # During Feature creation dataset samples might increase, we will select required samples again
                  predict_dataset = predict_dataset.select(
                      range(args.max_predict_samples)
                  )
  
      # Log a few random samples from the training set:
      for index in random.sample(range(len(train_dataset)), 3):
          logger.info(
              f"Sample {index} of the training set: {train_dataset[index]}."
          )
  
      # DataLoaders creation:
      if args.pad_to_max_length:
          # If padding was already done ot max length, we use the default data collator that will just convert everything
          # to tensors.
          data_collator = default_data_collator
      else:
          # Otherwise, `DataCollatorWithPadding` will apply dynamic padding for us (by padding to the maximum length of
          # the samples passed). When using mixed precision, we add `pad_to_multiple_of=8` to pad all tensors to multiple
          # of 8s, which will enable the use of Tensor Cores on NVIDIA hardware with compute capability >= 7.5 (Volta).
          data_collator = DataCollatorWithPadding(
              tokenizer,
              pad_to_multiple_of=(8 if accelerator.use_fp16 else None),
          )
  
      train_dataloader = DataLoader(
          train_dataset,
          shuffle=True,
          collate_fn=data_collator,
          batch_size=args.per_device_train_batch_size,
      )
  
      eval_dataset_for_model = eval_dataset.remove_columns(
          ["example_id", "offset_mapping"]
      )
      eval_dataloader = DataLoader(
          eval_dataset_for_model,
          collate_fn=data_collator,
          batch_size=args.per_device_eval_batch_size,
      )
  
      if args.do_predict:
          predict_dataset_for_model = predict_dataset.remove_columns(
              ["example_id", "offset_mapping"]
          )
          predict_dataloader = DataLoader(
              predict_dataset_for_model,
              collate_fn=data_collator,
              batch_size=args.per_device_eval_batch_size,
          )
  
      # Post-processing:
      def post_processing_function(
          examples, features, predictions, stage="eval"
      ):
          # Post-processing: we match the start logits and end logits to answers in the original context.
          (
              predictions,
              scores_diff_json,
          ) = postprocess_qa_predictions_with_beam_search(
              examples=examples,
              features=features,
              predictions=predictions,
              version_2_with_negative=args.version_2_with_negative,
              n_best_size=args.n_best_size,
              max_answer_length=args.max_answer_length,
              start_n_top=model.config.start_n_top,
              end_n_top=model.config.end_n_top,
              output_dir=args.output_dir,
              prefix=stage,
          )
          # Format the result to the format the metric expects.
          if args.version_2_with_negative:
              formatted_predictions = [
                  {
                      "id": k,
                      "prediction_text": v,
                      "no_answer_probability": scores_diff_json[k],
                  }
                  for k, v in predictions.items()
              ]
          else:
              formatted_predictions = [
                  {"id": k, "prediction_text": v}
                  for k, v in predictions.items()
              ]
  
          references = [
              {"id": ex["id"], "answers": ex[answer_column_name]}
              for ex in examples
          ]
          return EvalPrediction(
              predictions=formatted_predictions, label_ids=references
          )
  
      metric = evaluate.load(
          "squad_v2" if args.version_2_with_negative else "squad"
      )
  
      def create_and_fill_np_array(
          start_or_end_logits, dataset, max_len
      ):
          """
          Create and fill numpy array of size len_of_validation_data * max_length_of_output_tensor
  
          Args:
              start_or_end_logits(:obj:`tensor`):
                  This is the output predictions of the model. We can only enter either start or end logits.
              eval_dataset: Evaluation dataset
              max_len(:obj:`int`):
                  The maximum length of the output tensor. ( See the model.eval() part for more details )
          """
  
          step = 0
          # create a numpy array and fill it with -100.
          logits_concat = np.full(
              (len(dataset), max_len), -100, dtype=np.float32
          )
          # Now since we have create an array now we will populate it with the outputs gathered using accelerator.gather_for_metrics
          for i, output_logit in enumerate(
              start_or_end_logits
          ):  # populate columns
              # We have to fill it such that we have to take the whole tensor and replace it on the newly created array
              # And after every iteration we have to change the step
  
              batch_size = output_logit.shape[0]
              cols = output_logit.shape[1]
              if step + batch_size < len(dataset):
                  logits_concat[
                      step : step + batch_size, :cols
                  ] = output_logit
              else:
                  logits_concat[step:, :cols] = output_logit[
                      : len(dataset) - step
                  ]
  
              step += batch_size
  
          return logits_concat
  
      # Optimizer
      # Split weights in two groups, one with weight decay and the other not.
      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": args.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=args.learning_rate
      )
  
      # Scheduler and math around the number of training steps.
      overrode_max_train_steps = False
      num_update_steps_per_epoch = math.ceil(
          len(train_dataloader) / args.gradient_accumulation_steps
      )
      if args.max_train_steps is None:
          args.max_train_steps = (
              args.num_train_epochs * num_update_steps_per_epoch
          )
          overrode_max_train_steps = True
  
      lr_scheduler = get_scheduler(
          name=args.lr_scheduler_type,
          optimizer=optimizer,
          num_warmup_steps=args.num_warmup_steps
          * args.gradient_accumulation_steps,
          num_training_steps=args.max_train_steps
          * args.gradient_accumulation_steps,
      )
  
      # Prepare everything with our `accelerator`.
      (
          model,
          optimizer,
          train_dataloader,
          eval_dataloader,
          lr_scheduler,
      ) = accelerator.prepare(
          model,
          optimizer,
          train_dataloader,
          eval_dataloader,
          lr_scheduler,
      )
  
      # We need to recalculate our total training steps as the size of the training dataloader may have changed.
      num_update_steps_per_epoch = math.ceil(
          len(train_dataloader) / args.gradient_accumulation_steps
      )
      if overrode_max_train_steps:
          args.max_train_steps = (
              args.num_train_epochs * num_update_steps_per_epoch
          )
      # Afterwards we recalculate our number of training epochs
      args.num_train_epochs = math.ceil(
          args.max_train_steps / num_update_steps_per_epoch
      )
  
      # Figure out how many steps we should save the Accelerator states
      checkpointing_steps = args.checkpointing_steps
      if (
          checkpointing_steps is not None
          and checkpointing_steps.isdigit()
      ):
          checkpointing_steps = int(checkpointing_steps)
  
      # We need to initialize the trackers we use, and also store our configuration
      if args.with_tracking:
          experiment_config = vars(args)
          # TensorBoard cannot log Enums, need the raw value
          experiment_config["lr_scheduler_type"] = experiment_config[
              "lr_scheduler_type"
          ].value
          accelerator.init_trackers(
              "qa_beam_search_no_trainer", experiment_config
          )
  
      # Train!
      total_batch_size = (
          args.per_device_train_batch_size
          * accelerator.num_processes
          * args.gradient_accumulation_steps
      )
  
      logger.info("***** Running training *****")
      logger.info(f"  Num examples = {len(train_dataset)}")
      logger.info(f"  Num Epochs = {args.num_train_epochs}")
      logger.info(
          f"  Instantaneous batch size per device = {args.per_device_train_batch_size}"
      )
      logger.info(
          f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
      )
      logger.info(
          f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}"
      )
      logger.info(
          f"  Total optimization steps = {args.max_train_steps}"
      )
  
      # Only show the progress bar once on each machine.
      progress_bar = tqdm(
          range(args.max_train_steps),
          disable=not accelerator.is_local_main_process,
      )
      completed_steps = 0
      starting_epoch = 0
  
      # Potentially load in the weights and states from a previous save
      if args.resume_from_checkpoint:
          if (
              args.resume_from_checkpoint is not None
              or args.resume_from_checkpoint != ""
          ):
              accelerator.print(
                  f"Resumed from checkpoint: {args.resume_from_checkpoint}"
              )
              accelerator.load_state(args.resume_from_checkpoint)
              path = os.path.basename(args.resume_from_checkpoint)
          else:
              # Get the most recent checkpoint
              dirs = [
                  f.name for f in os.scandir(os.getcwd()) if f.is_dir()
              ]
              dirs.sort(key=os.path.getctime)
              path = dirs[
                  -1
              ]  # Sorts folders by date modified, most recent checkpoint is the last
          # Extract `epoch_{i}` or `step_{i}`
          training_difference = os.path.splitext(path)[0]
  
          if "epoch" in training_difference:
              starting_epoch = (
                  int(training_difference.replace("epoch_", "")) + 1
              )
              resume_step = None
          else:
              resume_step = int(
                  training_difference.replace("step_", "")
              )
              starting_epoch = resume_step // len(train_dataloader)
              resume_step -= starting_epoch * len(train_dataloader)
  
      for epoch in range(starting_epoch, args.num_train_epochs):
          model.train()
          if args.with_tracking:
              total_loss = 0
          for step, batch in enumerate(train_dataloader):
              # We need to skip steps until we reach the resumed step
              if (
                  args.resume_from_checkpoint
                  and epoch == starting_epoch
              ):
                  if resume_step is not None and step < resume_step:
                      completed_steps += 1
                      continue
  
              with accelerator.accumulate(model):
                  outputs = model(**batch)
                  loss = outputs.loss
                  # We keep track of the loss at each epoch
                  if args.with_tracking:
                      total_loss += loss.detach().float()
  
                  accelerator.backward(loss)
  
                  optimizer.step()
                  lr_scheduler.step()
                  optimizer.zero_grad()
  
              # Checks if the accelerator has performed an optimization step behind the scenes
              if accelerator.sync_gradients:
                  progress_bar.update(1)
                  completed_steps += 1
  
              if isinstance(checkpointing_steps, int):
                  if completed_steps % checkpointing_steps == 0:
                      accelerator.save_state(f"step_{completed_steps}")
  
              if completed_steps >= args.max_train_steps:
                  break
  
          if args.push_to_hub and epoch < args.num_train_epochs - 1:
              accelerator.wait_for_everyone()
              unwrapped_model = accelerator.unwrap_model(model)
              unwrapped_model.save_pretrained(
                  args.output_dir,
                  is_main_process=accelerator.is_main_process,
                  save_function=accelerator.save,
              )
              if accelerator.is_main_process:
                  tokenizer.save_pretrained(args.output_dir)
                  repo.push_to_hub(
                      commit_message=f"Training in progress epoch {epoch}",
                      blocking=False,
                      auto_lfs_prune=True,
                  )
  
      # intialize all lists to collect the batches
      all_start_top_log_probs = []
      all_start_top_index = []
      all_end_top_log_probs = []
      all_end_top_index = []
      all_cls_logits = []
  
      model.eval()
  
      for step, batch in enumerate(eval_dataloader):
          with torch.no_grad():
              outputs = model(**batch)
              start_top_log_probs = outputs.start_top_log_probs
              start_top_index = outputs.start_top_index
              end_top_log_probs = outputs.end_top_log_probs
              end_top_index = outputs.end_top_index
              cls_logits = outputs.cls_logits
  
              if (
                  not args.pad_to_max_length
              ):  # necessary to pad predictions and labels for being gathered
                  start_top_log_probs = (
                      accelerator.pad_across_processes(
                          start_top_log_probs, dim=1, pad_index=-100
                      )
                  )
                  start_top_index = accelerator.pad_across_processes(
                      start_top_index, dim=1, pad_index=-100
                  )
                  end_top_log_probs = accelerator.pad_across_processes(
                      end_top_log_probs, dim=1, pad_index=-100
                  )
                  end_top_index = accelerator.pad_across_processes(
                      end_top_index, dim=1, pad_index=-100
                  )
                  cls_logits = accelerator.pad_across_processes(
                      cls_logits, dim=1, pad_index=-100
                  )
  
              all_start_top_log_probs.append(
                  accelerator.gather_for_metrics(start_top_log_probs)
                  .cpu()
                  .numpy()
              )
              all_start_top_index.append(
                  accelerator.gather_for_metrics(start_top_index)
                  .cpu()
                  .numpy()
              )
              all_end_top_log_probs.append(
                  accelerator.gather_for_metrics(end_top_log_probs)
                  .cpu()
                  .numpy()
              )
              all_end_top_index.append(
                  accelerator.gather_for_metrics(end_top_index)
                  .cpu()
                  .numpy()
              )
              all_cls_logits.append(
                  accelerator.gather_for_metrics(cls_logits)
                  .cpu()
                  .numpy()
              )
  
      max_len = max(
          [x.shape[1] for x in all_end_top_log_probs]
      )  # Get the max_length of the tensor
  
      # concatenate all numpy arrays collected above
      start_top_log_probs_concat = create_and_fill_np_array(
          all_start_top_log_probs, eval_dataset, max_len
      )
      start_top_index_concat = create_and_fill_np_array(
          all_start_top_index, eval_dataset, max_len
      )
      end_top_log_probs_concat = create_and_fill_np_array(
          all_end_top_log_probs, eval_dataset, max_len
      )
      end_top_index_concat = create_and_fill_np_array(
          all_end_top_index, eval_dataset, max_len
      )
      cls_logits_concat = np.concatenate(all_cls_logits, axis=0)
  
      # delete the list of numpy arrays
      del start_top_log_probs
      del start_top_index
      del end_top_log_probs
      del end_top_index
      del cls_logits
  
      outputs_numpy = (
          start_top_log_probs_concat,
          start_top_index_concat,
          end_top_log_probs_concat,
          end_top_index_concat,
          cls_logits_concat,
      )
      prediction = post_processing_function(
          eval_examples, eval_dataset, outputs_numpy
      )
      eval_metric = metric.compute(
          predictions=prediction.predictions,
          references=prediction.label_ids,
      )
      logger.info(f"Evaluation metrics: {eval_metric}")
  
      if args.do_predict:
          # intialize all lists to collect the batches
  
          all_start_top_log_probs = []
          all_start_top_index = []
          all_end_top_log_probs = []
          all_end_top_index = []
          all_cls_logits = []
  
          model.eval()
  
          for step, batch in enumerate(predict_dataloader):
              with torch.no_grad():
                  outputs = model(**batch)
                  start_top_log_probs = outputs.start_top_log_probs
                  start_top_index = outputs.start_top_index
                  end_top_log_probs = outputs.end_top_log_probs
                  end_top_index = outputs.end_top_index
                  cls_logits = outputs.cls_logits
  
                  if (
                      not args.pad_to_max_length
                  ):  # necessary to pad predictions and labels for being gathered
                      start_top_log_probs = (
                          accelerator.pad_across_processes(
                              start_top_log_probs, dim=1, pad_index=-100
                          )
                      )
                      start_top_index = (
                          accelerator.pad_across_processes(
                              start_top_index, dim=1, pad_index=-100
                          )
                      )
                      end_top_log_probs = (
                          accelerator.pad_across_processes(
                              end_top_log_probs, dim=1, pad_index=-100
                          )
                      )
                      end_top_index = accelerator.pad_across_processes(
                          end_top_index, dim=1, pad_index=-100
                      )
                      cls_logits = accelerator.pad_across_processes(
                          cls_logits, dim=1, pad_index=-100
                      )
  
                  all_start_top_log_probs.append(
                      accelerator.gather_for_metrics(
                          start_top_log_probs
                      )
                      .cpu()
                      .numpy()
                  )
                  all_start_top_index.append(
                      accelerator.gather_for_metrics(start_top_index)
                      .cpu()
                      .numpy()
                  )
                  all_end_top_log_probs.append(
                      accelerator.gather_for_metrics(end_top_log_probs)
                      .cpu()
                      .numpy()
                  )
                  all_end_top_index.append(
                      accelerator.gather_for_metrics(end_top_index)
                      .cpu()
                      .numpy()
                  )
                  all_cls_logits.append(
                      accelerator.gather_for_metrics(cls_logits)
                      .cpu()
                      .numpy()
                  )
  
          max_len = max(
              [x.shape[1] for x in all_end_top_log_probs]
          )  # Get the max_length of the tensor
  
          # concatenate all numpy arrays collected above
          start_top_log_probs_concat = create_and_fill_np_array(
              all_start_top_log_probs, predict_dataset, max_len
          )
          start_top_index_concat = create_and_fill_np_array(
              all_start_top_index, predict_dataset, max_len
          )
          end_top_log_probs_concat = create_and_fill_np_array(
              all_end_top_log_probs, predict_dataset, max_len
          )
          end_top_index_concat = create_and_fill_np_array(
              all_end_top_index, predict_dataset, max_len
          )
          cls_logits_concat = np.concatenate(all_cls_logits, axis=0)
  
          # delete the list of numpy arrays
          del start_top_log_probs
          del start_top_index
          del end_top_log_probs
          del end_top_index
          del cls_logits
  
          outputs_numpy = (
              start_top_log_probs_concat,
              start_top_index_concat,
              end_top_log_probs_concat,
              end_top_index_concat,
              cls_logits_concat,
          )
  
          prediction = post_processing_function(
              predict_examples, predict_dataset, outputs_numpy
          )
          predict_metric = metric.compute(
              predictions=prediction.predictions,
              references=prediction.label_ids,
          )
          logger.info(f"Predict metrics: {predict_metric}")
  
      if args.with_tracking:
          log = {
              "squad_v2"
              if args.version_2_with_negative
              else "squad": eval_metric,
              "train_loss": total_loss,
              "epoch": epoch,
              "step": completed_steps,
          }
          if args.do_predict:
              log[
                  "squad_v2_predict"
                  if args.version_2_with_negative
                  else "squad_predict"
              ] = predict_metric
  
          accelerator.log(log)
  
      if args.checkpointing_steps == "epoch":
          accelerator.save_state(f"epoch_{epoch}")
  
      if args.output_dir is not None:
          accelerator.wait_for_everyone()
          unwrapped_model = accelerator.unwrap_model(model)
          unwrapped_model.save_pretrained(
              args.output_dir,
              is_main_process=accelerator.is_main_process,
              save_function=accelerator.save,
          )
          if accelerator.is_main_process:
              tokenizer.save_pretrained(args.output_dir)
              if args.push_to_hub:
                  repo.push_to_hub(
                      commit_message="End of training",
                      auto_lfs_prune=True,
                  )
  
              logger.info(json.dumps(eval_metric, indent=4))
              save_prefixed_metrics(eval_metric, args.output_dir)
  
imports: import argparse
  import json
  import logging
  import math
  import os
  import random
  from pathlib import Path
  import datasets
  import numpy
  import torch
  from datasets import load_dataset
  from torch.utils.data import DataLoader
  from tqdm.auto import tqdm
  import evaluate
  import transformers
  from accelerate import Accelerator
  from accelerate.logging import get_logger
  from accelerate.utils import set_seed
  from huggingface_hub import Repository
  from transformers import AdamW, DataCollatorWithPadding, EvalPrediction, SchedulerType, XLNetConfig, XLNetForQuestionAnswering, XLNetTokenizerFast, default_data_collator, get_scheduler
  from transformers.utils import check_min_version, get_full_repo_name, send_example_telemetry
  from transformers.utils.versions import require_version
  from utils_qa import postprocess_qa_predictions_with_beam_search
label_desc: Use image node_id 3 for calling a global function with example usage: main() and returns: tokenized_examples, tokenized_examples, EvalPrediction, logits_concat