code: def main():
      # region Argument parsing
      # See all possible arguments in src/transformers/training_args.py
      # or by passing the --help flag to this script.
      # We now keep distinct sets of args, for a cleaner separation of concerns.
  
      parser = HfArgumentParser(
          (ModelArguments, DataTrainingArguments, TrainingArguments)
      )
      if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
          # If we pass only one argument to the script and it's the path to a json file,
          # let's parse it to get our arguments.
          model_args, data_args, training_args = parser.parse_json_file(
              json_file=os.path.abspath(sys.argv[1])
          )
      else:
          (
              model_args,
              data_args,
              training_args,
          ) = parser.parse_args_into_dataclasses()
  
      if model_args.use_auth_token is not None:
          warnings.warn(
              "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead.",
              FutureWarning,
          )
          if model_args.token is not None:
              raise ValueError(
                  "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
              )
          model_args.token = model_args.use_auth_token
  
      # 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", model_args, data_args, framework="flax"
      )
      # endregion
  
      # region Logging
      # 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,
      )
      # Setup logging, we only want one process per machine to log things on the screen.
      logger.setLevel(
          logging.INFO if jax.process_index() == 0 else logging.ERROR
      )
      if jax.process_index() == 0:
          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()
      # endregion
  
      # Handle the repository creation
      if training_args.push_to_hub:
          # Retrieve of infer repo_name
          repo_name = training_args.hub_model_id
          if repo_name is None:
              repo_name = Path(training_args.output_dir).absolute().name
          # Create repo and retrieve repo_id
          repo_id = create_repo(
              repo_name, exist_ok=True, token=training_args.hub_token
          ).repo_id
          # Clone repo locally
          repo = Repository(
              training_args.output_dir,
              clone_from=repo_id,
              token=training_args.hub_token,
          )
  
      # region Load Data
      # 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 data_args.dataset_name is not None:
          # Downloading and loading a dataset from the hub.
          raw_datasets = load_dataset(
              data_args.dataset_name,
              data_args.dataset_config_name,
              cache_dir=model_args.cache_dir,
              token=model_args.token,
          )
      else:
          # Loading the dataset from local csv or json file.
          data_files = {}
          if data_args.train_file is not None:
              data_files["train"] = data_args.train_file
              extension = data_args.train_file.split(".")[-1]
  
          if data_args.validation_file is not None:
              data_files["validation"] = data_args.validation_file
              extension = data_args.validation_file.split(".")[-1]
          if data_args.test_file is not None:
              data_files["test"] = data_args.test_file
              extension = data_args.test_file.split(".")[-1]
          raw_datasets = load_dataset(
              extension,
              data_files=data_files,
              field="data",
              cache_dir=model_args.cache_dir,
              token=model_args.token,
          )
      # 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.
      # endregion
  
      # region Load pretrained model and tokenizer
      #
      # Load pretrained model and tokenizer
      config = AutoConfig.from_pretrained(
          model_args.config_name
          if model_args.config_name
          else model_args.model_name_or_path,
          cache_dir=model_args.cache_dir,
          revision=model_args.model_revision,
          token=model_args.token,
          trust_remote_code=model_args.trust_remote_code,
      )
      tokenizer = AutoTokenizer.from_pretrained(
          model_args.tokenizer_name
          if model_args.tokenizer_name
          else model_args.model_name_or_path,
          cache_dir=model_args.cache_dir,
          use_fast=True,
          revision=model_args.model_revision,
          token=model_args.token,
          trust_remote_code=model_args.trust_remote_code,
      )
      # endregion
  
      # region Tokenizer check: this script requires a fast tokenizer.
      if not isinstance(tokenizer, PreTrainedTokenizerFast):
          raise ValueError(
              "This example script only works for models that have a fast tokenizer. Checkout the big table of models at"
              " https://huggingface.co/transformers/index.html#supported-frameworks to find the model types that meet"
              " this requirement"
          )
      # endregion
  
      # region Preprocessing the datasets
      # Preprocessing is slightly different for training and evaluation.
      if training_args.do_train:
          column_names = raw_datasets["train"].column_names
      elif training_args.do_eval:
          column_names = raw_datasets["validation"].column_names
      else:
          column_names = raw_datasets["test"].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 data_args.max_seq_length > tokenizer.model_max_length:
          logger.warning(
              f"The max_seq_length passed ({data_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(
          data_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=data_args.doc_stride,
              return_overflowing_tokens=True,
              return_offsets_mapping=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")
  
          # Let's label those examples!
          tokenized_examples["start_positions"] = []
          tokenized_examples["end_positions"] = []
  
          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)
  
              # Grab the sequence corresponding to that example (to know what is the context and what is the question).
              sequence_ids = tokenized_examples.sequence_ids(i)
  
              # 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)
              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] != (
                      1 if pad_on_right else 0
                  ):
                      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] != (
                      1 if pad_on_right else 0
                  ):
                      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
                      )
                  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
                      )
  
          return tokenized_examples
  
      processed_raw_datasets = {}
      if training_args.do_train:
          if "train" not in raw_datasets:
              raise ValueError("--do_train requires a train dataset")
          train_dataset = raw_datasets["train"]
          if data_args.max_train_samples is not None:
              # We will select sample from whole data if agument is specified
              max_train_samples = min(
                  len(train_dataset), data_args.max_train_samples
              )
              train_dataset = train_dataset.select(
                  range(max_train_samples)
              )
          # Create train feature from dataset
          train_dataset = train_dataset.map(
              prepare_train_features,
              batched=True,
              num_proc=data_args.preprocessing_num_workers,
              remove_columns=column_names,
              load_from_cache_file=not data_args.overwrite_cache,
          )
          if data_args.max_train_samples is not None:
              # Number of samples might increase during Feature Creation, We select only specified max samples
              max_train_samples = min(
                  len(train_dataset), data_args.max_train_samples
              )
              train_dataset = train_dataset.select(
                  range(max_train_samples)
              )
          processed_raw_datasets["train"] = train_dataset
  
      # 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=data_args.doc_stride,
              return_overflowing_tokens=True,
              return_offsets_mapping=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"
          )
  
          # 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"] = []
  
          for i in range(len(tokenized_examples["input_ids"])):
              # Grab the sequence corresponding to that example (to know what is the context and what is the question).
              sequence_ids = tokenized_examples.sequence_ids(i)
              context_index = 1 if pad_on_right else 0
  
              # 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_index else None)
                  for k, o in enumerate(
                      tokenized_examples["offset_mapping"][i]
                  )
              ]
  
          return tokenized_examples
  
      if training_args.do_eval:
          if "validation" not in raw_datasets:
              raise ValueError(
                  "--do_eval requires a validation dataset"
              )
          eval_examples = raw_datasets["validation"]
          if data_args.max_eval_samples is not None:
              # We will select sample from whole data
              max_eval_samples = min(
                  len(eval_examples), data_args.max_eval_samples
              )
              eval_examples = eval_examples.select(
                  range(max_eval_samples)
              )
          # Validation Feature Creation
          eval_dataset = eval_examples.map(
              prepare_validation_features,
              batched=True,
              num_proc=data_args.preprocessing_num_workers,
              remove_columns=column_names,
              load_from_cache_file=not data_args.overwrite_cache,
          )
          if data_args.max_eval_samples is not None:
              # During Feature creation dataset samples might increase, we will select required samples again
              max_eval_samples = min(
                  len(eval_dataset), data_args.max_eval_samples
              )
              eval_dataset = eval_dataset.select(
                  range(max_eval_samples)
              )
          processed_raw_datasets["validation"] = eval_dataset
  
      if training_args.do_predict:
          if "test" not in raw_datasets:
              raise ValueError("--do_predict requires a test dataset")
          predict_examples = raw_datasets["test"]
          if data_args.max_predict_samples is not None:
              # We will select sample from whole data
              predict_examples = predict_examples.select(
                  range(data_args.max_predict_samples)
              )
          # Predict Feature Creation
          predict_dataset = predict_examples.map(
              prepare_validation_features,
              batched=True,
              num_proc=data_args.preprocessing_num_workers,
              remove_columns=column_names,
              load_from_cache_file=not data_args.overwrite_cache,
          )
          if data_args.max_predict_samples is not None:
              # During Feature creation dataset samples might increase, we will select required samples again
              max_predict_samples = min(
                  len(predict_dataset), data_args.max_predict_samples
              )
              predict_dataset = predict_dataset.select(
                  range(max_predict_samples)
              )
          processed_raw_datasets["test"] = predict_dataset
      # endregion
  
      # region Metrics and 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 = postprocess_qa_predictions(
              examples=examples,
              features=features,
              predictions=predictions,
              version_2_with_negative=data_args.version_2_with_negative,
              n_best_size=data_args.n_best_size,
              max_answer_length=data_args.max_answer_length,
              null_score_diff_threshold=data_args.null_score_diff_threshold,
              output_dir=training_args.output_dir,
              prefix=stage,
          )
          # Format the result to the format the metric expects.
          if data_args.version_2_with_negative:
              formatted_predictions = [
                  {
                      "id": k,
                      "prediction_text": v,
                      "no_answer_probability": 0.0,
                  }
                  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 data_args.version_2_with_negative else "squad",
          cache_dir=model_args.cache_dir,
      )
  
      def compute_metrics(p: EvalPrediction):
          return metric.compute(
              predictions=p.predictions, references=p.label_ids
          )
  
      # Create and fill numpy array of size len_of_validation_data * max_length_of_output_tensor
      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.float64
          )
          # Now since we have create an array now we will populate it with the outputs of the model.
          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
  
      # endregion
  
      # region Training steps and logging init
      train_dataset = processed_raw_datasets["train"]
      eval_dataset = processed_raw_datasets["validation"]
  
      # 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]}."
          )
  
      # Define a summary writer
      has_tensorboard = is_tensorboard_available()
      if has_tensorboard and jax.process_index() == 0:
          try:
              from flax.metrics.tensorboard import SummaryWriter
  
              summary_writer = SummaryWriter(training_args.output_dir)
              summary_writer.hparams(
                  {
                      **training_args.to_dict(),
                      **vars(model_args),
                      **vars(data_args),
                  }
              )
          except ImportError as ie:
              has_tensorboard = False
              logger.warning(
                  f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
              )
      else:
          logger.warning(
              "Unable to display metrics through TensorBoard because the package is not installed: "
              "Please run pip install tensorboard to enable."
          )
  
      def write_train_metric(
          summary_writer, train_metrics, train_time, step
      ):
          summary_writer.scalar("train_time", train_time, step)
  
          train_metrics = get_metrics(train_metrics)
          for key, vals in train_metrics.items():
              tag = f"train_{key}"
              for i, val in enumerate(vals):
                  summary_writer.scalar(
                      tag, val, step - len(vals) + i + 1
                  )
  
      def write_eval_metric(summary_writer, eval_metrics, step):
          for metric_name, value in eval_metrics.items():
              summary_writer.scalar(f"eval_{metric_name}", value, step)
  
      num_epochs = int(training_args.num_train_epochs)
      rng = jax.random.PRNGKey(training_args.seed)
      dropout_rngs = jax.random.split(rng, jax.local_device_count())
  
      train_batch_size = (
          int(training_args.per_device_train_batch_size)
          * jax.local_device_count()
      )
      per_device_eval_batch_size = int(
          training_args.per_device_eval_batch_size
      )
      eval_batch_size = (
          per_device_eval_batch_size * jax.local_device_count()
      )
      # endregion
  
      # region Load model
      model = FlaxAutoModelForQuestionAnswering.from_pretrained(
          model_args.model_name_or_path,
          config=config,
          cache_dir=model_args.cache_dir,
          revision=model_args.model_revision,
          token=model_args.token,
          trust_remote_code=model_args.trust_remote_code,
          seed=training_args.seed,
          dtype=getattr(jnp, model_args.dtype),
      )
  
      learning_rate_fn = create_learning_rate_fn(
          len(train_dataset),
          train_batch_size,
          training_args.num_train_epochs,
          training_args.warmup_steps,
          training_args.learning_rate,
      )
  
      state = create_train_state(
          model,
          learning_rate_fn,
          num_labels=max_seq_length,
          training_args=training_args,
      )
      # endregion
  
      # region Define train step functions
      def train_step(
          state: train_state.TrainState,
          batch: Dict[str, Array],
          dropout_rng: PRNGKey,
      ) -> Tuple[train_state.TrainState, float]:
          """Trains model with an optimizer (both in `state`) on `batch`, returning a pair `(new_state, loss)`."""
          dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
          start_positions = batch.pop("start_positions")
          end_positions = batch.pop("end_positions")
          targets = (start_positions, end_positions)
  
          def loss_fn(params):
              logits = state.apply_fn(
                  **batch,
                  params=params,
                  dropout_rng=dropout_rng,
                  train=True,
              )
              loss = state.loss_fn(logits, targets)
              return loss
  
          grad_fn = jax.value_and_grad(loss_fn)
          loss, grad = grad_fn(state.params)
          grad = jax.lax.pmean(grad, "batch")
          new_state = state.apply_gradients(grads=grad)
          metrics = jax.lax.pmean(
              {
                  "loss": loss,
                  "learning_rate": learning_rate_fn(state.step),
              },
              axis_name="batch",
          )
          return new_state, metrics, new_dropout_rng
  
      p_train_step = jax.pmap(
          train_step, axis_name="batch", donate_argnums=(0,)
      )
      # endregion
  
      # region Define eval step functions
      def eval_step(state, batch):
          logits = state.apply_fn(
              **batch, params=state.params, train=False
          )
          return state.logits_fn(logits)
  
      p_eval_step = jax.pmap(eval_step, axis_name="batch")
      # endregion
  
      # region Define train and eval loop
      logger.info(
          f"===== Starting training ({num_epochs} epochs) ====="
      )
      train_time = 0
  
      # make sure weights are replicated on each device
      state = replicate(state)
  
      train_time = 0
      step_per_epoch = len(train_dataset) // train_batch_size
      total_steps = step_per_epoch * num_epochs
      epochs = tqdm(
          range(num_epochs),
          desc=f"Epoch ... (1/{num_epochs})",
          position=0,
      )
      for epoch in epochs:
          train_start = time.time()
          train_metrics = []
  
          # Create sampling rng
          rng, input_rng = jax.random.split(rng)
  
          # train
          for step, batch in enumerate(
              tqdm(
                  train_data_collator(
                      input_rng, train_dataset, train_batch_size
                  ),
                  total=step_per_epoch,
                  desc="Training...",
                  position=1,
              ),
              1,
          ):
              state, train_metric, dropout_rngs = p_train_step(
                  state, batch, dropout_rngs
              )
              train_metrics.append(train_metric)
  
              cur_step = epoch * step_per_epoch + step
  
              if (
                  cur_step % training_args.logging_steps == 0
                  and cur_step > 0
              ):
                  # Save metrics
                  train_metric = unreplicate(train_metric)
                  train_time += time.time() - train_start
                  if has_tensorboard and jax.process_index() == 0:
                      write_train_metric(
                          summary_writer,
                          train_metrics,
                          train_time,
                          cur_step,
                      )
  
                  epochs.write(
                      f"Step... ({cur_step}/{total_steps} | Training Loss: {train_metric['loss']}, Learning Rate:"
                      f" {train_metric['learning_rate']})"
                  )
  
                  train_metrics = []
  
              if (
                  training_args.do_eval
                  and (
                      cur_step % training_args.eval_steps == 0
                      or cur_step % step_per_epoch == 0
                  )
                  and cur_step > 0
              ):
                  eval_metrics = {}
                  all_start_logits = []
                  all_end_logits = []
                  # evaluate
                  for batch in tqdm(
                      eval_data_collator(eval_dataset, eval_batch_size),
                      total=math.ceil(
                          len(eval_dataset) / eval_batch_size
                      ),
                      desc="Evaluating ...",
                      position=2,
                  ):
                      _ = batch.pop("example_id")
                      _ = batch.pop("offset_mapping")
                      predictions = pad_shard_unpad(p_eval_step)(
                          state,
                          batch,
                          min_device_batch=per_device_eval_batch_size,
                      )
                      start_logits = np.array(predictions[0])
                      end_logits = np.array(predictions[1])
                      all_start_logits.append(start_logits)
                      all_end_logits.append(end_logits)
  
                  max_len = max(
                      [x.shape[1] for x in all_start_logits]
                  )  # Get the max_length of the tensor
  
                  # concatenate the numpy array
                  start_logits_concat = create_and_fill_np_array(
                      all_start_logits, eval_dataset, max_len
                  )
                  end_logits_concat = create_and_fill_np_array(
                      all_end_logits, eval_dataset, max_len
                  )
  
                  # delete the list of numpy arrays
                  del all_start_logits
                  del all_end_logits
                  outputs_numpy = (
                      start_logits_concat,
                      end_logits_concat,
                  )
                  prediction = post_processing_function(
                      eval_examples, eval_dataset, outputs_numpy
                  )
                  eval_metrics = compute_metrics(prediction)
  
                  logger.info(
                      f"Step... ({cur_step}/{total_steps} | Evaluation metrics: {eval_metrics})"
                  )
  
                  if has_tensorboard and jax.process_index() == 0:
                      write_eval_metric(
                          summary_writer, eval_metrics, cur_step
                      )
  
              if (
                  cur_step % training_args.save_steps == 0
                  and cur_step > 0
              ) or (cur_step == total_steps):
                  # save checkpoint after each epoch and push checkpoint to the hub
                  if jax.process_index() == 0:
                      params = jax.device_get(unreplicate(state.params))
                      model.save_pretrained(
                          training_args.output_dir, params=params
                      )
                      tokenizer.save_pretrained(
                          training_args.output_dir
                      )
                      if training_args.push_to_hub:
                          repo.push_to_hub(
                              commit_message=f"Saving weights and logs of step {cur_step}",
                              blocking=False,
                          )
          epochs.desc = f"Epoch ... {epoch + 1}/{num_epochs}"
      # endregion
  
      # Eval after training
      if training_args.do_eval:
          eval_metrics = {}
          all_start_logits = []
          all_end_logits = []
  
          eval_loader = eval_data_collator(
              eval_dataset, eval_batch_size
          )
          for batch in tqdm(
              eval_loader,
              total=math.ceil(len(eval_dataset) / eval_batch_size),
              desc="Evaluating ...",
              position=2,
          ):
              _ = batch.pop("example_id")
              _ = batch.pop("offset_mapping")
              predictions = pad_shard_unpad(p_eval_step)(
                  state,
                  batch,
                  min_device_batch=per_device_eval_batch_size,
              )
              start_logits = np.array(predictions[0])
              end_logits = np.array(predictions[1])
              all_start_logits.append(start_logits)
              all_end_logits.append(end_logits)
  
          max_len = max(
              [x.shape[1] for x in all_start_logits]
          )  # Get the max_length of the tensor
  
          # concatenate the numpy array
          start_logits_concat = create_and_fill_np_array(
              all_start_logits, eval_dataset, max_len
          )
          end_logits_concat = create_and_fill_np_array(
              all_end_logits, eval_dataset, max_len
          )
  
          # delete the list of numpy arrays
          del all_start_logits
          del all_end_logits
          outputs_numpy = (start_logits_concat, end_logits_concat)
          prediction = post_processing_function(
              eval_examples, eval_dataset, outputs_numpy
          )
          eval_metrics = compute_metrics(prediction)
  
          if jax.process_index() == 0:
              eval_metrics = {
                  f"eval_{metric_name}": value
                  for metric_name, value in eval_metrics.items()
              }
              path = os.path.join(
                  training_args.output_dir, "eval_results.json"
              )
              with open(path, "w") as f:
                  json.dump(eval_metrics, f, indent=4, sort_keys=True)
  
imports: import json
  import logging
  import math
  import os
  import random
  import sys
  import time
  import warnings
  from dataclasses import asdict, dataclass, field
  from enum import Enum
  from pathlib import Path
  from typing import Any, Callable, Dict, Optional, Tuple
  import datasets
  import evaluate
  import jax
  import jax.numpy
  import numpy
  import optax
  from datasets import load_dataset
  from flax import struct, traverse_util
  from flax.jax_utils import pad_shard_unpad, replicate, unreplicate
  from flax.training import train_state
  from flax.training.common_utils import get_metrics, onehot, shard
  from huggingface_hub import Repository, create_repo
  from tqdm import tqdm
  from utils_qa import postprocess_qa_predictions
  import transformers
  from transformers import AutoConfig, AutoTokenizer, EvalPrediction, FlaxAutoModelForQuestionAnswering, HfArgumentParser, PreTrainedTokenizerFast, is_tensorboard_available
  from transformers.utils import check_min_version, send_example_telemetry
label_desc: Use image node_id 5 for calling a global function with example usage: main() and returns: tokenized_examples, tokenized_examples, EvalPrediction, metric, logits_concat, new_state,  metrics,  new_dropout_rng, state, loss