code: def masking_function(
      text: str,
      tokenizer: TokenizerType,
      mask_prob: float,
      random_replace_prob: float,
      unmask_replace_prob: float,
      max_length: int,
  ) -> Tuple[List[int], List[int]]:
      """Given a text string, randomly mask wordpieces for Bert MLM
      training.
  
      Args:
          text (str):
              The input text
          tokenizer (TokenizerType):
              The tokenizer for tokenization
          mask_prob (float):
              What fraction of tokens to mask
          random_replace_prob (float):
              Of the masked tokens, how many should be replaced with
              random tokens (improves performance)
          unmask_replace_prob (float):
              Of the masked tokens, how many should be replaced with
              the original token (improves performance)
          max_length (int):
              The maximum sequence length to consider. Note that for
              Bert style models, this is a function of the number of
              positional embeddings you learn
  
      Returns:
          Tuple[List[int], List[int]]:
              The masked token ids (based on the tokenizer passed),
              and the output labels (padded with `tokenizer.pad_token_id`)
      """
      # Note: By default, encode does add the BOS and EOS token
      # Disabling that behaviour to make this more clear
      tokenized_ids = (
          [tokenizer.bos_token_id]
          + tokenizer.encode(
              text,
              add_special_tokens=False,
              truncation=True,
              max_length=max_length - 2,
          )
          + [tokenizer.eos_token_id]
      )
      seq_len = len(tokenized_ids)
      tokenized_ids = np.array(tokenized_ids)
      subword_mask = np.full(len(tokenized_ids), False)
  
      # Masking the BOS and EOS token leads to slightly worse performance
      low = 1
      high = len(subword_mask) - 1
      mask_choices = np.arange(low, high)
      num_subwords_to_mask = max(
          int((mask_prob * (high - low)) + np.random.rand()), 1
      )
      subword_mask[
          np.random.choice(
              mask_choices, num_subwords_to_mask, replace=False
          )
      ] = True
  
      # Create the labels first
      labels = np.full(seq_len, tokenizer.pad_token_id)
      labels[subword_mask] = tokenized_ids[subword_mask]
  
      tokenized_ids[subword_mask] = tokenizer.mask_token_id
  
      # Now of the masked tokens, choose how many to replace with random and how many to unmask
      rand_or_unmask_prob = random_replace_prob + unmask_replace_prob
      if rand_or_unmask_prob > 0:
          rand_or_unmask = subword_mask & (
              np.random.rand(len(tokenized_ids)) < rand_or_unmask_prob
          )
          if random_replace_prob == 0:
              unmask = rand_or_unmask
              rand_mask = None
          elif unmask_replace_prob == 0:
              unmask = None
              rand_mask = rand_or_unmask
          else:
              unmask_prob = unmask_replace_prob / rand_or_unmask_prob
              decision = (
                  np.random.rand(len(tokenized_ids)) < unmask_prob
              )
              unmask = rand_or_unmask & decision
              rand_mask = rand_or_unmask & (~decision)
          if unmask is not None:
              tokenized_ids[unmask] = labels[unmask]
          if rand_mask is not None:
              weights = np.ones(tokenizer.vocab_size)
              weights[tokenizer.all_special_ids] = 0
              probs = weights / weights.sum()
              num_rand = rand_mask.sum()
              tokenized_ids[rand_mask] = np.random.choice(
                  tokenizer.vocab_size, num_rand, p=probs
              )
      return tokenized_ids.tolist(), labels.tolist()
  
imports: import datetime
  import json
  import pathlib
  import re
  import string
  from functools import partial
  from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, TypeVar, Union
  import random
  import datasets
  import fire
  import loguru
  import numpy
  import pytz
  import sh
  import torch
  import torch.nn
  from torch.utils.data import DataLoader, Dataset
  from torch.utils.tensorboard import SummaryWriter
  from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
  from transformers.models.roberta import RobertaConfig, RobertaModel
  from transformers.models.roberta.modeling_roberta import RobertaLMHead, RobertaPreTrainedModel
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 2 for calling a global function with example usage: masking_function(text, tokenizer, mask_prob, random_replace_prob, unmask_replace_prob, max_length) and returns: tokenized_ids,  labels