code: def reduce_next_sentence_label_dimension(batch):
      """Change next_sentence_labels's shape from (-1, 1) to (-1,).
  
      Args:
        batch: A dictionary mapping keys to arrays.
  
      Returns:
        Updated batch.
      """
  
      batch["next_sentence_labels"] = batch["next_sentence_labels"][
          :, 0
      ]
      return batch
  
file_path: scenic/scenic/projects/baselines/bert/datasets/bert_wikibooks_dataset.py
imports: import functools
  from typing import Optional
  from absl import logging
  from flax import jax_utils
  import jax
  import jax.numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  import tensorflow
  from official.nlp.data import pretrain_dataloader
label_desc: Use image node_id 1 for calling a global function with example usage: reduce_next_sentence_label_dimension(batch) and returns: batch