class_docstr: Data collator used for BART denoising language modeling. The code is largely copied from
  `<https://github.com/morganmcg1/rotobart/blob/main/data_collator.py#L223>`__.
  For more information on how BART denoising language modeling works, one can take a look
  at the `official paper <https://arxiv.org/pdf/1910.13461.pdf>`__
  or the `official code for preprocessing <https://github.com/facebookresearch/fairseq/blob/main/fairseq/data/denoising_dataset.py>`__ .
  Args:
      tokenizer (:class:`~transformers.PreTrainedTokenizer` or :class:`~transformers.PreTrainedTokenizerFast`):
          The tokenizer used for encoding the data
      mask_ratio (:obj:`float`):
          The probability with which to (randomly) mask tokens in the input
      poisson_lambda (:obj:`float`):
          Mean parameter of Poisson distribution used to generate span-lengths to be masked
      permute_sentence_ratio (:obj:`float`):
          Ratio of sentences to be permuted in each document
      decoder_start_token_id: (:obj:`int):
          The decoder start token id of the model
docstr: Sampling text spans with span lengths drawn from a Poisson distribution and masking them.
docstr_tok: ["Sampling","text","spans","with","span","lengths","drawn","from","a","Poisson","distribution","and","masking","them","."]
code: def span_mask_tokens(self, input_ids, labels, do_permute):
      """
      Sampling text spans with span lengths drawn from a Poisson distribution and masking them.
      """
      special_tokens_mask_labels = [
          self.tokenizer.get_special_tokens_mask(
              val, already_has_special_tokens=True
          )
          for val in labels.tolist()
      ]
      special_tokens_mask_inputs = [
          self.tokenizer.get_special_tokens_mask(
              val, already_has_special_tokens=True
          )
          for val in input_ids.tolist()
      ]
      special_tokens_mask_labels = np.array(
          special_tokens_mask_labels, dtype=bool
      )
      special_tokens_mask_inputs = np.array(
          special_tokens_mask_inputs, dtype=bool
      )
  
      # determine how many tokens we need to mask in total
      is_token_mask = (
          ~(input_ids == self.tokenizer.pad_token_id)
          & ~special_tokens_mask_inputs
      )
      num_tokens_to_mask = int(
          math.ceil(is_token_mask.astype(float).sum() * self.mask_ratio)
      )
      if num_tokens_to_mask == 0:
          return input_ids, labels
  
      # generate a sufficient number of span lengths
      span_lengths = np.random.poisson(
          lam=self.poisson_lambda, size=(num_tokens_to_mask,)
      )
      while np.cumsum(span_lengths, 0)[-1] < num_tokens_to_mask:
          span_lengths = np.concatenate(
              [
                  span_lengths,
                  np.random.poisson(
                      lam=self.poisson_lambda,
                      size=(num_tokens_to_mask,),
                  ),
              ]
          )
  
      # remove all spans of length 0
      # note that BART inserts additional mask tokens where length == 0,
      # which we do not implement for now as it adds additional complexity
      span_lengths = span_lengths[span_lengths > 0]
  
      # trim to about num_tokens_to_mask tokens
      cutoff_idx = (
          np.argmin(
              np.abs(np.cumsum(span_lengths, 0) - num_tokens_to_mask)
          )
          + 1
      )
      span_lengths = span_lengths[:cutoff_idx]
  
      # randomly choose starting positions for masking
      token_indices = np.argwhere(is_token_mask == 1)
      span_starts = np.random.permutation(token_indices.shape[0])[
          : span_lengths.shape[0]
      ]
      # prepare mask
      masked_indices = np.array(token_indices[span_starts])
      mask = np.full_like(input_ids, fill_value=False)
  
      # mask starting positions
      for mi in masked_indices:
          mask[tuple(mi)] = True
      span_lengths -= 1
  
      # fill up spans
      max_index = input_ids.shape[1] - 1
      remaining = (span_lengths > 0) & (
          masked_indices[:, 1] < max_index
      )
      while np.any(remaining):
          masked_indices[remaining, 1] += 1
          for mi in masked_indices:
              mask[tuple(mi)] = True
          span_lengths -= 1
          remaining = (span_lengths > 0) & (
              masked_indices[:, 1] < max_index
          )
  
      # place the mask tokens
      mask[np.where(special_tokens_mask_inputs)] = False
      input_ids[np.where(mask)] = self.tokenizer.mask_token_id
      if not do_permute:
          labels[np.where(mask == 0)] = -100
      else:
          labels[np.where(special_tokens_mask_labels)] = -100
  
      # remove mask tokens that are not starts of spans
      to_remove = (mask == 1) & np.roll((mask == 1), 1, 1)
      new_input_ids = np.full_like(
          input_ids, fill_value=self.tokenizer.pad_token_id
      )
      for i, example in enumerate(input_ids):
          new_example = example[~to_remove[i]]
          new_input_ids[i, : new_example.shape[0]] = new_example
  
      return new_input_ids, labels
  
imports: import json
  import logging
  import math
  import os
  import sys
  import time
  from dataclasses import asdict, dataclass, field
  from enum import Enum
  from itertools import chain
  from pathlib import Path
  from typing import Dict, List, Optional
  import nltk
  import numpy
  from datasets import load_dataset
  from tqdm import tqdm
  import flax
  import jax
  import jax.numpy
  import optax
  from flax import jax_utils, traverse_util
  from flax.jax_utils import pad_shard_unpad
  from flax.training import train_state
  from flax.training.common_utils import get_metrics, onehot, shard
  from huggingface_hub import Repository
  from transformers import CONFIG_MAPPING, FLAX_MODEL_FOR_MASKED_LM_MAPPING, AutoTokenizer, BartConfig, BatchEncoding, FlaxBartForConditionalGeneration, HfArgumentParser, PreTrainedTokenizerBase, is_tensorboard_available, set_seed
  from transformers.models.bart.modeling_flax_bart import shift_tokens_right
  from transformers.utils import get_full_repo_name, send_example_telemetry
label_desc: Use image node_id 4 for calling the FlaxDataCollatorForBartDenoisingLM obj's underlying member method code with example usage: obj.span_mask_tokens(input_ids, labels, do_permute) and returns: new_input_ids,  labels, input_ids,  labels