class_docstr: Data collator used for language modeling. Inputs are dynamically padded to the maximum length of a batch if they
  are not all of the same length.
  
  Args:
      tokenizer (:class:`~transformers.PreTrainedTokenizer` or :class:`~transformers.PreTrainedTokenizerFast`):
          The tokenizer used for encoding the data.
      mlm (:obj:`bool`, `optional`, defaults to :obj:`True`):
          Whether or not to use masked language modeling. If set to :obj:`False`, the labels are the same as the
          inputs with the padding tokens ignored (by setting them to -100). Otherwise, the labels are -100 for
          non-masked tokens and the value to predict for the masked token.
      mlm_probability (:obj:`float`, `optional`, defaults to 0.15):
          The probability with which to (randomly) mask tokens in the input, when :obj:`mlm` is set to :obj:`True`.
  
  .. note::
  
      For best performance, this data collator should be used with a dataset having items that are dictionaries or
      BatchEncoding, with the :obj:`"special_tokens_mask"` key, as returned by a
      :class:`~transformers.PreTrainedTokenizer` or a :class:`~transformers.PreTrainedTokenizerFast` with the
      argument :obj:`return_special_tokens_mask=True`.
class_docstr_tok: ["Data","collator","used","for","language","modeling",".","Inputs","are","dynamically","padded","to","the","maximum","length","of","a","batch","if","they","are","not","all","of","the","same","length",".","Args",":","tokenizer","(",":","class",":","`","~transformers.PreTrainedTokenizer","`","or",":","class",":","`","~transformers.PreTrainedTokenizerFast","`",")",":","The","tokenizer","used","for","encoding","the","data",".","mlm","(",":","obj",":","`","bool","`",",","`","optional","`",",","defaults","to",":","obj",":","`","True","`",")",":","Whether","or","not","to","use","masked","language","modeling",".","If","set","to",":","obj",":","`","False","`",",","the","labels","are","the","same","as","the","inputs","with","the","padding","tokens","ignored","(","by","setting","them","to","-100",")",".","Otherwise",",","the","labels","are","-100","for","non-masked","tokens","and","the","value","to","predict","for","the","masked","token",".","mlm_probability","(",":","obj",":","`","float","`",",","`","optional","`",",","defaults","to","0.15",")",":","The","probability","with","which","to","(","randomly",")","mask","tokens","in","the","input",",","when",":","obj",":","`","mlm","`","is","set","to",":","obj",":","`","True","`",".","..","note",":",":","For","best","performance",",","this","data","collator","should","be","used","with","a","dataset","having","items","that","are","dictionaries","or","BatchEncoding",",","with","the",":","obj",":","`","``","special_tokens_mask","''","`","key",",","as","returned","by","a",":","class",":","`","~transformers.PreTrainedTokenizer","`","or","a",":","class",":","`","~transformers.PreTrainedTokenizerFast","`","with","the","argument",":","obj",":","`","return_special_tokens_mask=True","`","."]
docstr: Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
docstr_tok: ["Prepare","masked","tokens","inputs\/labels","for","masked","language","modeling",":","80","%","MASK",",","10","%","random",",","10","%","original","."]
code: def mask_tokens(
      self,
      inputs: np.ndarray,
      special_tokens_mask: Optional[np.ndarray],
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """
      Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
      """
      labels = inputs.copy()
      # We sample a few tokens in each sequence for MLM training (with probability `self.mlm_probability`)
      probability_matrix = np.full(labels.shape, self.mlm_probability)
      special_tokens_mask = special_tokens_mask.astype("bool")
  
      probability_matrix[special_tokens_mask] = 0.0
      masked_indices = np.random.binomial(1, probability_matrix).astype(
          "bool"
      )
      labels[
          ~masked_indices
      ] = -100  # We only compute loss on masked tokens
  
      # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
      indices_replaced = (
          np.random.binomial(1, np.full(labels.shape, 0.8)).astype(
              "bool"
          )
          & masked_indices
      )
      inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(
          self.tokenizer.mask_token
      )
  
      # 10% of the time, we replace masked input tokens with random word
      indices_random = np.random.binomial(
          1, np.full(labels.shape, 0.5)
      ).astype("bool")
      indices_random &= masked_indices & ~indices_replaced
  
      random_words = np.random.randint(
          self.tokenizer.vocab_size, size=labels.shape, dtype="i4"
      )
      inputs[indices_random] = random_words[indices_random]
  
      # The rest of the time (10% of the time) we keep the masked input tokens unchanged
      return inputs, labels
  
imports: import logging
  import os
  import sys
  from dataclasses import dataclass, field
  from pathlib import Path
  from typing import Dict, List, Optional, Tuple
  import jax
  import jax.numpy
  import numpy
  from datasets import load_dataset
  from flax import jax_utils
  from flax.optim import Adam
  from flax.training import common_utils
  from flax.training.common_utils import get_metrics
  from jax.nn import log_softmax
  from modeling_flax_performer import FlaxPerformerForMaskedLM
  from tqdm import tqdm
  from transformers import MODEL_FOR_MASKED_LM_MAPPING, AutoTokenizer, BertConfig, FlaxBertForMaskedLM, HfArgumentParser, PreTrainedTokenizerBase, TensorType, TrainingArguments, is_tensorboard_available, set_seed
label_desc: Use image node_id 3 for calling the FlaxDataCollatorForLanguageModeling obj's underlying member method code with example usage: obj.mask_tokens(inputs, special_tokens_mask) and returns: inputs,  labels