code: def create_masked_lm_predictions(
      tokens,
      vocab_id_list,
      vocab_id_to_token_dict,
      masked_lm_prob,
      cls_id,
      sep_id,
      mask_id,
      max_predictions_per_seq,
      np_rng,
      max_ngrams=3,
      do_whole_word_mask=True,
      favor_longer_ngram=False,
      do_permutation=False,
  ):
      """Creates the predictions for the masked LM objective.
      Note: Tokens here are vocab ids and not text tokens."""
  
      cand_indexes = []
      # Note(mingdachen): We create a list for recording if the piece is
      # the starting piece of current token, where 1 means true, so that
      # on-the-fly whole word masking is possible.
      token_boundary = [0] * len(tokens)
  
      for i, token in enumerate(tokens):
          if token == cls_id or token == sep_id:
              token_boundary[i] = 1
              continue
          # Whole Word Masking means that if we mask all of the wordpieces
          # corresponding to an original word.
          #
          # Note that Whole Word Masking does *not* change the training code
          # at all -- we still predict each WordPiece independently, softmaxed
          # over the entire vocabulary.
          if (
              do_whole_word_mask
              and len(cand_indexes) >= 1
              and not is_start_piece(vocab_id_to_token_dict[token])
          ):
              cand_indexes[-1].append(i)
          else:
              cand_indexes.append([i])
              if is_start_piece(vocab_id_to_token_dict[token]):
                  token_boundary[i] = 1
  
      output_tokens = list(tokens)
  
      masked_lm_positions = []
      masked_lm_labels = []
  
      if masked_lm_prob == 0:
          return (
              output_tokens,
              masked_lm_positions,
              masked_lm_labels,
              token_boundary,
          )
  
      num_to_predict = min(
          max_predictions_per_seq,
          max(1, int(round(len(tokens) * masked_lm_prob))),
      )
  
      # Note(mingdachen):
      # By default, we set the probabilities to favor shorter ngram sequences.
      ngrams = np.arange(1, max_ngrams + 1, dtype=np.int64)
      pvals = 1.0 / np.arange(1, max_ngrams + 1)
      pvals /= pvals.sum(keepdims=True)
  
      if favor_longer_ngram:
          pvals = pvals[::-1]
  
      ngram_indexes = []
      for idx in range(len(cand_indexes)):
          ngram_index = []
          for n in ngrams:
              ngram_index.append(cand_indexes[idx : idx + n])
          ngram_indexes.append(ngram_index)
  
      np_rng.shuffle(ngram_indexes)
  
      masked_lms = []
      covered_indexes = set()
      for cand_index_set in ngram_indexes:
          if len(masked_lms) >= num_to_predict:
              break
          if not cand_index_set:
              continue
          # Note(mingdachen):
          # Skip current piece if they are covered in lm masking or previous ngrams.
          for index_set in cand_index_set[0]:
              for index in index_set:
                  if index in covered_indexes:
                      continue
  
          n = np_rng.choice(
              ngrams[: len(cand_index_set)],
              p=pvals[: len(cand_index_set)]
              / pvals[: len(cand_index_set)].sum(keepdims=True),
          )
          index_set = sum(cand_index_set[n - 1], [])
          n -= 1
          # Note(mingdachen):
          # Repeatedly looking for a candidate that does not exceed the
          # maximum number of predictions by trying shorter ngrams.
          while len(masked_lms) + len(index_set) > num_to_predict:
              if n == 0:
                  break
              index_set = sum(cand_index_set[n - 1], [])
              n -= 1
          # If adding a whole-word mask would exceed the maximum number of
          # predictions, then just skip this candidate.
          if len(masked_lms) + len(index_set) > num_to_predict:
              continue
          is_any_index_covered = False
          for index in index_set:
              if index in covered_indexes:
                  is_any_index_covered = True
                  break
          if is_any_index_covered:
              continue
          for index in index_set:
              covered_indexes.add(index)
  
              masked_token = None
              # 80% of the time, replace with [MASK]
              if np_rng.random() < 0.8:
                  masked_token = mask_id
              else:
                  # 10% of the time, keep original
                  if np_rng.random() < 0.5:
                      masked_token = tokens[index]
                  # 10% of the time, replace with random word
                  else:
                      masked_token = vocab_id_list[
                          np_rng.randint(0, len(vocab_id_list))
                      ]
  
              output_tokens[index] = masked_token
  
              masked_lms.append(
                  MaskedLmInstance(index=index, label=tokens[index])
              )
      assert len(masked_lms) <= num_to_predict
  
      np_rng.shuffle(ngram_indexes)
  
      select_indexes = set()
      if do_permutation:
          for cand_index_set in ngram_indexes:
              if len(select_indexes) >= num_to_predict:
                  break
              if not cand_index_set:
                  continue
              # Note(mingdachen):
              # Skip current piece if they are covered in lm masking or previous ngrams.
              for index_set in cand_index_set[0]:
                  for index in index_set:
                      if (
                          index in covered_indexes
                          or index in select_indexes
                      ):
                          continue
  
              n = np.random.choice(
                  ngrams[: len(cand_index_set)],
                  p=pvals[: len(cand_index_set)]
                  / pvals[: len(cand_index_set)].sum(keepdims=True),
              )
              index_set = sum(cand_index_set[n - 1], [])
              n -= 1
  
              while (
                  len(select_indexes) + len(index_set) > num_to_predict
              ):
                  if n == 0:
                      break
                  index_set = sum(cand_index_set[n - 1], [])
                  n -= 1
              # If adding a whole-word mask would exceed the maximum number of
              # predictions, then just skip this candidate.
              if len(select_indexes) + len(index_set) > num_to_predict:
                  continue
              is_any_index_covered = False
              for index in index_set:
                  if (
                      index in covered_indexes
                      or index in select_indexes
                  ):
                      is_any_index_covered = True
                      break
              if is_any_index_covered:
                  continue
              for index in index_set:
                  select_indexes.add(index)
          assert len(select_indexes) <= num_to_predict
  
          select_indexes = sorted(select_indexes)
          permute_indexes = list(select_indexes)
          np_rng.shuffle(permute_indexes)
          orig_token = list(output_tokens)
  
          for src_i, tgt_i in zip(select_indexes, permute_indexes):
              output_tokens[src_i] = orig_token[tgt_i]
              masked_lms.append(
                  MaskedLmInstance(index=src_i, label=orig_token[src_i])
              )
  
      masked_lms = sorted(masked_lms, key=lambda x: x.index)
  
      for p in masked_lms:
          masked_lm_positions.append(p.index)
          masked_lm_labels.append(p.label)
  
      return (
          output_tokens,
          masked_lm_positions,
          masked_lm_labels,
          token_boundary,
      )
  
file_path: ColossalAI/examples/tutorial/sequence_parallel/data/datasets/dataset_utils.py
imports: import collections
  import math
  import time
  import numpy
  from colossalai.logging import get_dist_logger
  from .blendable_dataset import BlendableDataset
  from .indexed_dataset import make_dataset
label_desc: Use image node_id 7 for calling a global function with example usage: create_masked_lm_predictions(tokens, vocab_id_list, vocab_id_to_token_dict, masked_lm_prob, cls_id, sep_id, mask_id, max_predictions_per_seq, np_rng, max_ngrams, do_whole_word_mask, favor_longer_ngram, do_permutation) and returns: output_tokens,  masked_lm_positions,  masked_lm_labels,  token_boundary, output_tokens,  masked_lm_positions,  masked_lm_labels,  token_boundary